1use crate::{
2 prompts::PromptBuilder, slash_command::SlashCommandLine, AssistantPanel, InitialInsertion,
3 InlineAssistId, InlineAssistant, MessageId, MessageStatus,
4};
5use anyhow::{anyhow, Context as _, Result};
6use assistant_slash_command::{
7 SlashCommandOutput, SlashCommandOutputSection, SlashCommandRegistry,
8};
9use client::{self, proto, telemetry::Telemetry};
10use clock::ReplicaId;
11use collections::{HashMap, HashSet};
12use editor::Editor;
13use fs::{Fs, RemoveOptions};
14use futures::{
15 future::{self, Shared},
16 FutureExt, StreamExt,
17};
18use gpui::{
19 AppContext, Context as _, EventEmitter, Model, ModelContext, Subscription, Task, UpdateGlobal,
20 View, WeakView,
21};
22use language::{
23 AnchorRangeExt, Bias, Buffer, BufferSnapshot, LanguageRegistry, OffsetRangeExt, ParseStatus,
24 Point, ToOffset,
25};
26use language_model::{
27 LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, LanguageModelTool,
28 Role,
29};
30use open_ai::Model as OpenAiModel;
31use paths::contexts_dir;
32use project::Project;
33use schemars::JsonSchema;
34use serde::{Deserialize, Serialize};
35use std::{
36 cmp,
37 fmt::Debug,
38 iter, mem,
39 ops::Range,
40 path::{Path, PathBuf},
41 sync::Arc,
42 time::{Duration, Instant},
43};
44use telemetry_events::AssistantKind;
45use ui::{SharedString, WindowContext};
46use util::{post_inc, ResultExt, TryFutureExt};
47use uuid::Uuid;
48use workspace::Workspace;
49
50#[derive(Clone, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
51pub struct ContextId(String);
52
53impl ContextId {
54 pub fn new() -> Self {
55 Self(Uuid::new_v4().to_string())
56 }
57
58 pub fn from_proto(id: String) -> Self {
59 Self(id)
60 }
61
62 pub fn to_proto(&self) -> String {
63 self.0.clone()
64 }
65}
66
67#[derive(Clone, Debug)]
68pub enum ContextOperation {
69 InsertMessage {
70 anchor: MessageAnchor,
71 metadata: MessageMetadata,
72 version: clock::Global,
73 },
74 UpdateMessage {
75 message_id: MessageId,
76 metadata: MessageMetadata,
77 version: clock::Global,
78 },
79 UpdateSummary {
80 summary: ContextSummary,
81 version: clock::Global,
82 },
83 SlashCommandFinished {
84 id: SlashCommandId,
85 output_range: Range<language::Anchor>,
86 sections: Vec<SlashCommandOutputSection<language::Anchor>>,
87 version: clock::Global,
88 },
89 BufferOperation(language::Operation),
90}
91
92impl ContextOperation {
93 pub fn from_proto(op: proto::ContextOperation) -> Result<Self> {
94 match op.variant.context("invalid variant")? {
95 proto::context_operation::Variant::InsertMessage(insert) => {
96 let message = insert.message.context("invalid message")?;
97 let id = MessageId(language::proto::deserialize_timestamp(
98 message.id.context("invalid id")?,
99 ));
100 Ok(Self::InsertMessage {
101 anchor: MessageAnchor {
102 id,
103 start: language::proto::deserialize_anchor(
104 message.start.context("invalid anchor")?,
105 )
106 .context("invalid anchor")?,
107 },
108 metadata: MessageMetadata {
109 role: Role::from_proto(message.role),
110 status: MessageStatus::from_proto(
111 message.status.context("invalid status")?,
112 ),
113 timestamp: id.0,
114 },
115 version: language::proto::deserialize_version(&insert.version),
116 })
117 }
118 proto::context_operation::Variant::UpdateMessage(update) => Ok(Self::UpdateMessage {
119 message_id: MessageId(language::proto::deserialize_timestamp(
120 update.message_id.context("invalid message id")?,
121 )),
122 metadata: MessageMetadata {
123 role: Role::from_proto(update.role),
124 status: MessageStatus::from_proto(update.status.context("invalid status")?),
125 timestamp: language::proto::deserialize_timestamp(
126 update.timestamp.context("invalid timestamp")?,
127 ),
128 },
129 version: language::proto::deserialize_version(&update.version),
130 }),
131 proto::context_operation::Variant::UpdateSummary(update) => Ok(Self::UpdateSummary {
132 summary: ContextSummary {
133 text: update.summary,
134 done: update.done,
135 timestamp: language::proto::deserialize_timestamp(
136 update.timestamp.context("invalid timestamp")?,
137 ),
138 },
139 version: language::proto::deserialize_version(&update.version),
140 }),
141 proto::context_operation::Variant::SlashCommandFinished(finished) => {
142 Ok(Self::SlashCommandFinished {
143 id: SlashCommandId(language::proto::deserialize_timestamp(
144 finished.id.context("invalid id")?,
145 )),
146 output_range: language::proto::deserialize_anchor_range(
147 finished.output_range.context("invalid range")?,
148 )?,
149 sections: finished
150 .sections
151 .into_iter()
152 .map(|section| {
153 Ok(SlashCommandOutputSection {
154 range: language::proto::deserialize_anchor_range(
155 section.range.context("invalid range")?,
156 )?,
157 icon: section.icon_name.parse()?,
158 label: section.label.into(),
159 })
160 })
161 .collect::<Result<Vec<_>>>()?,
162 version: language::proto::deserialize_version(&finished.version),
163 })
164 }
165 proto::context_operation::Variant::BufferOperation(op) => Ok(Self::BufferOperation(
166 language::proto::deserialize_operation(
167 op.operation.context("invalid buffer operation")?,
168 )?,
169 )),
170 }
171 }
172
173 pub fn to_proto(&self) -> proto::ContextOperation {
174 match self {
175 Self::InsertMessage {
176 anchor,
177 metadata,
178 version,
179 } => proto::ContextOperation {
180 variant: Some(proto::context_operation::Variant::InsertMessage(
181 proto::context_operation::InsertMessage {
182 message: Some(proto::ContextMessage {
183 id: Some(language::proto::serialize_timestamp(anchor.id.0)),
184 start: Some(language::proto::serialize_anchor(&anchor.start)),
185 role: metadata.role.to_proto() as i32,
186 status: Some(metadata.status.to_proto()),
187 }),
188 version: language::proto::serialize_version(version),
189 },
190 )),
191 },
192 Self::UpdateMessage {
193 message_id,
194 metadata,
195 version,
196 } => proto::ContextOperation {
197 variant: Some(proto::context_operation::Variant::UpdateMessage(
198 proto::context_operation::UpdateMessage {
199 message_id: Some(language::proto::serialize_timestamp(message_id.0)),
200 role: metadata.role.to_proto() as i32,
201 status: Some(metadata.status.to_proto()),
202 timestamp: Some(language::proto::serialize_timestamp(metadata.timestamp)),
203 version: language::proto::serialize_version(version),
204 },
205 )),
206 },
207 Self::UpdateSummary { summary, version } => proto::ContextOperation {
208 variant: Some(proto::context_operation::Variant::UpdateSummary(
209 proto::context_operation::UpdateSummary {
210 summary: summary.text.clone(),
211 done: summary.done,
212 timestamp: Some(language::proto::serialize_timestamp(summary.timestamp)),
213 version: language::proto::serialize_version(version),
214 },
215 )),
216 },
217 Self::SlashCommandFinished {
218 id,
219 output_range,
220 sections,
221 version,
222 } => proto::ContextOperation {
223 variant: Some(proto::context_operation::Variant::SlashCommandFinished(
224 proto::context_operation::SlashCommandFinished {
225 id: Some(language::proto::serialize_timestamp(id.0)),
226 output_range: Some(language::proto::serialize_anchor_range(
227 output_range.clone(),
228 )),
229 sections: sections
230 .iter()
231 .map(|section| {
232 let icon_name: &'static str = section.icon.into();
233 proto::SlashCommandOutputSection {
234 range: Some(language::proto::serialize_anchor_range(
235 section.range.clone(),
236 )),
237 icon_name: icon_name.to_string(),
238 label: section.label.to_string(),
239 }
240 })
241 .collect(),
242 version: language::proto::serialize_version(version),
243 },
244 )),
245 },
246 Self::BufferOperation(operation) => proto::ContextOperation {
247 variant: Some(proto::context_operation::Variant::BufferOperation(
248 proto::context_operation::BufferOperation {
249 operation: Some(language::proto::serialize_operation(operation)),
250 },
251 )),
252 },
253 }
254 }
255
256 fn timestamp(&self) -> clock::Lamport {
257 match self {
258 Self::InsertMessage { anchor, .. } => anchor.id.0,
259 Self::UpdateMessage { metadata, .. } => metadata.timestamp,
260 Self::UpdateSummary { summary, .. } => summary.timestamp,
261 Self::SlashCommandFinished { id, .. } => id.0,
262 Self::BufferOperation(_) => {
263 panic!("reading the timestamp of a buffer operation is not supported")
264 }
265 }
266 }
267
268 /// Returns the current version of the context operation.
269 pub fn version(&self) -> &clock::Global {
270 match self {
271 Self::InsertMessage { version, .. }
272 | Self::UpdateMessage { version, .. }
273 | Self::UpdateSummary { version, .. }
274 | Self::SlashCommandFinished { version, .. } => version,
275 Self::BufferOperation(_) => {
276 panic!("reading the version of a buffer operation is not supported")
277 }
278 }
279 }
280}
281
282#[derive(Debug, Clone)]
283pub enum ContextEvent {
284 AssistError(String),
285 MessagesEdited,
286 SummaryChanged,
287 WorkflowStepsChanged,
288 StreamedCompletion,
289 PendingSlashCommandsUpdated {
290 removed: Vec<Range<language::Anchor>>,
291 updated: Vec<PendingSlashCommand>,
292 },
293 SlashCommandFinished {
294 output_range: Range<language::Anchor>,
295 sections: Vec<SlashCommandOutputSection<language::Anchor>>,
296 run_commands_in_output: bool,
297 },
298 Operation(ContextOperation),
299}
300
301#[derive(Clone, Default, Debug)]
302pub struct ContextSummary {
303 pub text: String,
304 done: bool,
305 timestamp: clock::Lamport,
306}
307
308#[derive(Clone, Debug, Eq, PartialEq)]
309pub struct MessageAnchor {
310 pub id: MessageId,
311 pub start: language::Anchor,
312}
313
314#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
315pub struct MessageMetadata {
316 pub role: Role,
317 status: MessageStatus,
318 timestamp: clock::Lamport,
319}
320
321#[derive(Clone, Debug, PartialEq, Eq)]
322pub struct Message {
323 pub offset_range: Range<usize>,
324 pub index_range: Range<usize>,
325 pub id: MessageId,
326 pub anchor: language::Anchor,
327 pub role: Role,
328 pub status: MessageStatus,
329}
330
331impl Message {
332 fn to_request_message(&self, buffer: &Buffer) -> LanguageModelRequestMessage {
333 LanguageModelRequestMessage {
334 role: self.role,
335 content: buffer.text_for_range(self.offset_range.clone()).collect(),
336 }
337 }
338}
339
340struct PendingCompletion {
341 id: usize,
342 _task: Task<()>,
343}
344
345#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
346pub struct SlashCommandId(clock::Lamport);
347
348#[derive(Debug)]
349pub struct WorkflowStep {
350 pub tagged_range: Range<language::Anchor>,
351 pub edit_suggestions: WorkflowStepEditSuggestions,
352}
353
354#[derive(Clone, Debug, Eq, PartialEq)]
355pub struct ResolvedWorkflowStepEditSuggestions {
356 pub title: String,
357 pub edit_suggestions: HashMap<Model<Buffer>, Vec<EditSuggestionGroup>>,
358}
359
360pub enum WorkflowStepEditSuggestions {
361 Pending(Task<Option<()>>),
362 Resolved(ResolvedWorkflowStepEditSuggestions),
363}
364
365impl WorkflowStepEditSuggestions {
366 pub fn as_resolved(&self) -> Option<&ResolvedWorkflowStepEditSuggestions> {
367 match self {
368 WorkflowStepEditSuggestions::Resolved(suggestions) => Some(suggestions),
369 WorkflowStepEditSuggestions::Pending(_) => None,
370 }
371 }
372}
373
374#[derive(Clone, Debug, Eq, PartialEq)]
375pub struct EditSuggestionGroup {
376 pub context_range: Range<language::Anchor>,
377 pub suggestions: Vec<EditSuggestion>,
378}
379
380#[derive(Clone, Debug, Eq, PartialEq)]
381pub enum EditSuggestion {
382 Update {
383 range: Range<language::Anchor>,
384 description: String,
385 },
386 CreateFile {
387 description: String,
388 },
389 InsertSiblingBefore {
390 position: language::Anchor,
391 description: String,
392 },
393 InsertSiblingAfter {
394 position: language::Anchor,
395 description: String,
396 },
397 PrependChild {
398 position: language::Anchor,
399 description: String,
400 },
401 AppendChild {
402 position: language::Anchor,
403 description: String,
404 },
405 Delete {
406 range: Range<language::Anchor>,
407 },
408}
409
410impl EditSuggestion {
411 pub fn range(&self) -> Range<language::Anchor> {
412 match self {
413 EditSuggestion::Update { range, .. } => range.clone(),
414 EditSuggestion::CreateFile { .. } => language::Anchor::MIN..language::Anchor::MAX,
415 EditSuggestion::InsertSiblingBefore { position, .. }
416 | EditSuggestion::InsertSiblingAfter { position, .. }
417 | EditSuggestion::PrependChild { position, .. }
418 | EditSuggestion::AppendChild { position, .. } => *position..*position,
419 EditSuggestion::Delete { range } => range.clone(),
420 }
421 }
422
423 pub fn description(&self) -> Option<&str> {
424 match self {
425 EditSuggestion::Update { description, .. }
426 | EditSuggestion::CreateFile { description }
427 | EditSuggestion::InsertSiblingBefore { description, .. }
428 | EditSuggestion::InsertSiblingAfter { description, .. }
429 | EditSuggestion::PrependChild { description, .. }
430 | EditSuggestion::AppendChild { description, .. } => Some(description),
431 EditSuggestion::Delete { .. } => None,
432 }
433 }
434
435 fn description_mut(&mut self) -> Option<&mut String> {
436 match self {
437 EditSuggestion::Update { description, .. }
438 | EditSuggestion::CreateFile { description }
439 | EditSuggestion::InsertSiblingBefore { description, .. }
440 | EditSuggestion::InsertSiblingAfter { description, .. }
441 | EditSuggestion::PrependChild { description, .. }
442 | EditSuggestion::AppendChild { description, .. } => Some(description),
443 EditSuggestion::Delete { .. } => None,
444 }
445 }
446
447 fn try_merge(&mut self, other: &Self, buffer: &BufferSnapshot) -> bool {
448 let range = self.range();
449 let other_range = other.range();
450
451 // Don't merge if we don't contain the other suggestion.
452 if range.start.cmp(&other_range.start, buffer).is_gt()
453 || range.end.cmp(&other_range.end, buffer).is_lt()
454 {
455 return false;
456 }
457
458 if let Some(description) = self.description_mut() {
459 if let Some(other_description) = other.description() {
460 description.push('\n');
461 description.push_str(other_description);
462 }
463 }
464 true
465 }
466
467 pub fn show(
468 &self,
469 editor: &View<Editor>,
470 excerpt_id: editor::ExcerptId,
471 workspace: &WeakView<Workspace>,
472 assistant_panel: &View<AssistantPanel>,
473 cx: &mut WindowContext,
474 ) -> Option<InlineAssistId> {
475 let mut initial_transaction_id = None;
476 let initial_prompt;
477 let suggestion_range;
478 let buffer = editor.read(cx).buffer().clone();
479 let snapshot = buffer.read(cx).snapshot(cx);
480
481 match self {
482 EditSuggestion::Update { range, description } => {
483 initial_prompt = description.clone();
484 suggestion_range = snapshot.anchor_in_excerpt(excerpt_id, range.start)?
485 ..snapshot.anchor_in_excerpt(excerpt_id, range.end)?;
486 }
487 EditSuggestion::CreateFile { description } => {
488 initial_prompt = description.clone();
489 suggestion_range = editor::Anchor::min()..editor::Anchor::min();
490 }
491 EditSuggestion::InsertSiblingBefore {
492 position,
493 description,
494 } => {
495 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
496 initial_prompt = description.clone();
497 suggestion_range = buffer.update(cx, |buffer, cx| {
498 buffer.start_transaction(cx);
499 let line_start = buffer.insert_empty_line(position, true, true, cx);
500 initial_transaction_id = buffer.end_transaction(cx);
501
502 let line_start = buffer.read(cx).anchor_before(line_start);
503 line_start..line_start
504 });
505 }
506 EditSuggestion::InsertSiblingAfter {
507 position,
508 description,
509 } => {
510 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
511 initial_prompt = description.clone();
512 suggestion_range = buffer.update(cx, |buffer, cx| {
513 buffer.start_transaction(cx);
514 let line_start = buffer.insert_empty_line(position, true, true, cx);
515 initial_transaction_id = buffer.end_transaction(cx);
516
517 let line_start = buffer.read(cx).anchor_before(line_start);
518 line_start..line_start
519 });
520 }
521 EditSuggestion::PrependChild {
522 position,
523 description,
524 } => {
525 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
526 initial_prompt = description.clone();
527 suggestion_range = buffer.update(cx, |buffer, cx| {
528 buffer.start_transaction(cx);
529 let line_start = buffer.insert_empty_line(position, false, true, cx);
530 initial_transaction_id = buffer.end_transaction(cx);
531
532 let line_start = buffer.read(cx).anchor_before(line_start);
533 line_start..line_start
534 });
535 }
536 EditSuggestion::AppendChild {
537 position,
538 description,
539 } => {
540 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
541 initial_prompt = description.clone();
542 suggestion_range = buffer.update(cx, |buffer, cx| {
543 buffer.start_transaction(cx);
544 let line_start = buffer.insert_empty_line(position, true, false, cx);
545 initial_transaction_id = buffer.end_transaction(cx);
546
547 let line_start = buffer.read(cx).anchor_before(line_start);
548 line_start..line_start
549 });
550 }
551 EditSuggestion::Delete { range } => {
552 initial_prompt = "Delete".to_string();
553 suggestion_range = snapshot.anchor_in_excerpt(excerpt_id, range.start)?
554 ..snapshot.anchor_in_excerpt(excerpt_id, range.end)?;
555 }
556 }
557
558 InlineAssistant::update_global(cx, |inline_assistant, cx| {
559 Some(inline_assistant.suggest_assist(
560 editor,
561 suggestion_range,
562 initial_prompt,
563 initial_transaction_id,
564 Some(workspace.clone()),
565 Some(assistant_panel),
566 cx,
567 ))
568 })
569 }
570}
571
572impl Debug for WorkflowStepEditSuggestions {
573 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
574 match self {
575 WorkflowStepEditSuggestions::Pending(_) => write!(f, "EditStepOperations::Pending"),
576 WorkflowStepEditSuggestions::Resolved(ResolvedWorkflowStepEditSuggestions {
577 title,
578 edit_suggestions,
579 }) => f
580 .debug_struct("EditStepOperations::Parsed")
581 .field("title", title)
582 .field("edit_suggestions", edit_suggestions)
583 .finish(),
584 }
585 }
586}
587
588pub struct Context {
589 id: ContextId,
590 timestamp: clock::Lamport,
591 version: clock::Global,
592 pending_ops: Vec<ContextOperation>,
593 operations: Vec<ContextOperation>,
594 buffer: Model<Buffer>,
595 pending_slash_commands: Vec<PendingSlashCommand>,
596 edits_since_last_slash_command_parse: language::Subscription,
597 finished_slash_commands: HashSet<SlashCommandId>,
598 slash_command_output_sections: Vec<SlashCommandOutputSection<language::Anchor>>,
599 message_anchors: Vec<MessageAnchor>,
600 messages_metadata: HashMap<MessageId, MessageMetadata>,
601 summary: Option<ContextSummary>,
602 pending_summary: Task<Option<()>>,
603 completion_count: usize,
604 pending_completions: Vec<PendingCompletion>,
605 token_count: Option<usize>,
606 pending_token_count: Task<Option<()>>,
607 pending_save: Task<Result<()>>,
608 path: Option<PathBuf>,
609 _subscriptions: Vec<Subscription>,
610 telemetry: Option<Arc<Telemetry>>,
611 language_registry: Arc<LanguageRegistry>,
612 workflow_steps: Vec<WorkflowStep>,
613 project: Option<Model<Project>>,
614 prompt_builder: Arc<PromptBuilder>,
615}
616
617impl EventEmitter<ContextEvent> for Context {}
618
619impl Context {
620 pub fn local(
621 language_registry: Arc<LanguageRegistry>,
622 project: Option<Model<Project>>,
623 telemetry: Option<Arc<Telemetry>>,
624 prompt_builder: Arc<PromptBuilder>,
625 cx: &mut ModelContext<Self>,
626 ) -> Self {
627 Self::new(
628 ContextId::new(),
629 ReplicaId::default(),
630 language::Capability::ReadWrite,
631 language_registry,
632 prompt_builder,
633 project,
634 telemetry,
635 cx,
636 )
637 }
638
639 #[allow(clippy::too_many_arguments)]
640 pub fn new(
641 id: ContextId,
642 replica_id: ReplicaId,
643 capability: language::Capability,
644 language_registry: Arc<LanguageRegistry>,
645 prompt_builder: Arc<PromptBuilder>,
646 project: Option<Model<Project>>,
647 telemetry: Option<Arc<Telemetry>>,
648 cx: &mut ModelContext<Self>,
649 ) -> Self {
650 let buffer = cx.new_model(|_cx| {
651 let mut buffer = Buffer::remote(
652 language::BufferId::new(1).unwrap(),
653 replica_id,
654 capability,
655 "",
656 );
657 buffer.set_language_registry(language_registry.clone());
658 buffer
659 });
660 let edits_since_last_slash_command_parse =
661 buffer.update(cx, |buffer, _| buffer.subscribe());
662 let mut this = Self {
663 id,
664 timestamp: clock::Lamport::new(replica_id),
665 version: clock::Global::new(),
666 pending_ops: Vec::new(),
667 operations: Vec::new(),
668 message_anchors: Default::default(),
669 messages_metadata: Default::default(),
670 pending_slash_commands: Vec::new(),
671 finished_slash_commands: HashSet::default(),
672 slash_command_output_sections: Vec::new(),
673 edits_since_last_slash_command_parse,
674 summary: None,
675 pending_summary: Task::ready(None),
676 completion_count: Default::default(),
677 pending_completions: Default::default(),
678 token_count: None,
679 pending_token_count: Task::ready(None),
680 _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
681 pending_save: Task::ready(Ok(())),
682 path: None,
683 buffer,
684 telemetry,
685 project,
686 language_registry,
687 workflow_steps: Vec::new(),
688 prompt_builder,
689 };
690
691 let first_message_id = MessageId(clock::Lamport {
692 replica_id: 0,
693 value: 0,
694 });
695 let message = MessageAnchor {
696 id: first_message_id,
697 start: language::Anchor::MIN,
698 };
699 this.messages_metadata.insert(
700 first_message_id,
701 MessageMetadata {
702 role: Role::User,
703 status: MessageStatus::Done,
704 timestamp: first_message_id.0,
705 },
706 );
707 this.message_anchors.push(message);
708
709 this.set_language(cx);
710 this.count_remaining_tokens(cx);
711 this
712 }
713
714 fn serialize(&self, cx: &AppContext) -> SavedContext {
715 let buffer = self.buffer.read(cx);
716 SavedContext {
717 id: Some(self.id.clone()),
718 zed: "context".into(),
719 version: SavedContext::VERSION.into(),
720 text: buffer.text(),
721 messages: self
722 .messages(cx)
723 .map(|message| SavedMessage {
724 id: message.id,
725 start: message.offset_range.start,
726 metadata: self.messages_metadata[&message.id].clone(),
727 })
728 .collect(),
729 summary: self
730 .summary
731 .as_ref()
732 .map(|summary| summary.text.clone())
733 .unwrap_or_default(),
734 slash_command_output_sections: self
735 .slash_command_output_sections
736 .iter()
737 .filter_map(|section| {
738 let range = section.range.to_offset(buffer);
739 if section.range.start.is_valid(buffer) && !range.is_empty() {
740 Some(assistant_slash_command::SlashCommandOutputSection {
741 range,
742 icon: section.icon,
743 label: section.label.clone(),
744 })
745 } else {
746 None
747 }
748 })
749 .collect(),
750 }
751 }
752
753 #[allow(clippy::too_many_arguments)]
754 pub fn deserialize(
755 saved_context: SavedContext,
756 path: PathBuf,
757 language_registry: Arc<LanguageRegistry>,
758 prompt_builder: Arc<PromptBuilder>,
759 project: Option<Model<Project>>,
760 telemetry: Option<Arc<Telemetry>>,
761 cx: &mut ModelContext<Self>,
762 ) -> Self {
763 let id = saved_context.id.clone().unwrap_or_else(|| ContextId::new());
764 let mut this = Self::new(
765 id,
766 ReplicaId::default(),
767 language::Capability::ReadWrite,
768 language_registry,
769 prompt_builder,
770 project,
771 telemetry,
772 cx,
773 );
774 this.path = Some(path);
775 this.buffer.update(cx, |buffer, cx| {
776 buffer.set_text(saved_context.text.as_str(), cx)
777 });
778 let operations = saved_context.into_ops(&this.buffer, cx);
779 this.apply_ops(operations, cx).unwrap();
780 this
781 }
782
783 pub fn id(&self) -> &ContextId {
784 &self.id
785 }
786
787 pub fn replica_id(&self) -> ReplicaId {
788 self.timestamp.replica_id
789 }
790
791 pub fn version(&self, cx: &AppContext) -> ContextVersion {
792 ContextVersion {
793 context: self.version.clone(),
794 buffer: self.buffer.read(cx).version(),
795 }
796 }
797
798 pub fn set_capability(
799 &mut self,
800 capability: language::Capability,
801 cx: &mut ModelContext<Self>,
802 ) {
803 self.buffer
804 .update(cx, |buffer, cx| buffer.set_capability(capability, cx));
805 }
806
807 fn next_timestamp(&mut self) -> clock::Lamport {
808 let timestamp = self.timestamp.tick();
809 self.version.observe(timestamp);
810 timestamp
811 }
812
813 pub fn serialize_ops(
814 &self,
815 since: &ContextVersion,
816 cx: &AppContext,
817 ) -> Task<Vec<proto::ContextOperation>> {
818 let buffer_ops = self
819 .buffer
820 .read(cx)
821 .serialize_ops(Some(since.buffer.clone()), cx);
822
823 let mut context_ops = self
824 .operations
825 .iter()
826 .filter(|op| !since.context.observed(op.timestamp()))
827 .cloned()
828 .collect::<Vec<_>>();
829 context_ops.extend(self.pending_ops.iter().cloned());
830
831 cx.background_executor().spawn(async move {
832 let buffer_ops = buffer_ops.await;
833 context_ops.sort_unstable_by_key(|op| op.timestamp());
834 buffer_ops
835 .into_iter()
836 .map(|op| proto::ContextOperation {
837 variant: Some(proto::context_operation::Variant::BufferOperation(
838 proto::context_operation::BufferOperation {
839 operation: Some(op),
840 },
841 )),
842 })
843 .chain(context_ops.into_iter().map(|op| op.to_proto()))
844 .collect()
845 })
846 }
847
848 pub fn apply_ops(
849 &mut self,
850 ops: impl IntoIterator<Item = ContextOperation>,
851 cx: &mut ModelContext<Self>,
852 ) -> Result<()> {
853 let mut buffer_ops = Vec::new();
854 for op in ops {
855 match op {
856 ContextOperation::BufferOperation(buffer_op) => buffer_ops.push(buffer_op),
857 op @ _ => self.pending_ops.push(op),
858 }
859 }
860 self.buffer
861 .update(cx, |buffer, cx| buffer.apply_ops(buffer_ops, cx))?;
862 self.flush_ops(cx);
863
864 Ok(())
865 }
866
867 fn flush_ops(&mut self, cx: &mut ModelContext<Context>) {
868 let mut messages_changed = false;
869 let mut summary_changed = false;
870
871 self.pending_ops.sort_unstable_by_key(|op| op.timestamp());
872 for op in mem::take(&mut self.pending_ops) {
873 if !self.can_apply_op(&op, cx) {
874 self.pending_ops.push(op);
875 continue;
876 }
877
878 let timestamp = op.timestamp();
879 match op.clone() {
880 ContextOperation::InsertMessage {
881 anchor, metadata, ..
882 } => {
883 if self.messages_metadata.contains_key(&anchor.id) {
884 // We already applied this operation.
885 } else {
886 self.insert_message(anchor, metadata, cx);
887 messages_changed = true;
888 }
889 }
890 ContextOperation::UpdateMessage {
891 message_id,
892 metadata: new_metadata,
893 ..
894 } => {
895 let metadata = self.messages_metadata.get_mut(&message_id).unwrap();
896 if new_metadata.timestamp > metadata.timestamp {
897 *metadata = new_metadata;
898 messages_changed = true;
899 }
900 }
901 ContextOperation::UpdateSummary {
902 summary: new_summary,
903 ..
904 } => {
905 if self
906 .summary
907 .as_ref()
908 .map_or(true, |summary| new_summary.timestamp > summary.timestamp)
909 {
910 self.summary = Some(new_summary);
911 summary_changed = true;
912 }
913 }
914 ContextOperation::SlashCommandFinished {
915 id,
916 output_range,
917 sections,
918 ..
919 } => {
920 if self.finished_slash_commands.insert(id) {
921 let buffer = self.buffer.read(cx);
922 self.slash_command_output_sections
923 .extend(sections.iter().cloned());
924 self.slash_command_output_sections
925 .sort_by(|a, b| a.range.cmp(&b.range, buffer));
926 cx.emit(ContextEvent::SlashCommandFinished {
927 output_range,
928 sections,
929 run_commands_in_output: false,
930 });
931 }
932 }
933 ContextOperation::BufferOperation(_) => unreachable!(),
934 }
935
936 self.version.observe(timestamp);
937 self.timestamp.observe(timestamp);
938 self.operations.push(op);
939 }
940
941 if messages_changed {
942 cx.emit(ContextEvent::MessagesEdited);
943 cx.notify();
944 }
945
946 if summary_changed {
947 cx.emit(ContextEvent::SummaryChanged);
948 cx.notify();
949 }
950 }
951
952 fn can_apply_op(&self, op: &ContextOperation, cx: &AppContext) -> bool {
953 if !self.version.observed_all(op.version()) {
954 return false;
955 }
956
957 match op {
958 ContextOperation::InsertMessage { anchor, .. } => self
959 .buffer
960 .read(cx)
961 .version
962 .observed(anchor.start.timestamp),
963 ContextOperation::UpdateMessage { message_id, .. } => {
964 self.messages_metadata.contains_key(message_id)
965 }
966 ContextOperation::UpdateSummary { .. } => true,
967 ContextOperation::SlashCommandFinished {
968 output_range,
969 sections,
970 ..
971 } => {
972 let version = &self.buffer.read(cx).version;
973 sections
974 .iter()
975 .map(|section| §ion.range)
976 .chain([output_range])
977 .all(|range| {
978 let observed_start = range.start == language::Anchor::MIN
979 || range.start == language::Anchor::MAX
980 || version.observed(range.start.timestamp);
981 let observed_end = range.end == language::Anchor::MIN
982 || range.end == language::Anchor::MAX
983 || version.observed(range.end.timestamp);
984 observed_start && observed_end
985 })
986 }
987 ContextOperation::BufferOperation(_) => {
988 panic!("buffer operations should always be applied")
989 }
990 }
991 }
992
993 fn push_op(&mut self, op: ContextOperation, cx: &mut ModelContext<Self>) {
994 self.operations.push(op.clone());
995 cx.emit(ContextEvent::Operation(op));
996 }
997
998 pub fn buffer(&self) -> &Model<Buffer> {
999 &self.buffer
1000 }
1001
1002 pub fn path(&self) -> Option<&Path> {
1003 self.path.as_deref()
1004 }
1005
1006 pub fn summary(&self) -> Option<&ContextSummary> {
1007 self.summary.as_ref()
1008 }
1009
1010 pub fn workflow_steps(&self) -> &[WorkflowStep] {
1011 &self.workflow_steps
1012 }
1013
1014 pub fn workflow_step_for_range(&self, range: Range<language::Anchor>) -> Option<&WorkflowStep> {
1015 self.workflow_steps
1016 .iter()
1017 .find(|step| step.tagged_range == range)
1018 }
1019
1020 pub fn pending_slash_commands(&self) -> &[PendingSlashCommand] {
1021 &self.pending_slash_commands
1022 }
1023
1024 pub fn slash_command_output_sections(&self) -> &[SlashCommandOutputSection<language::Anchor>] {
1025 &self.slash_command_output_sections
1026 }
1027
1028 fn set_language(&mut self, cx: &mut ModelContext<Self>) {
1029 let markdown = self.language_registry.language_for_name("Markdown");
1030 cx.spawn(|this, mut cx| async move {
1031 let markdown = markdown.await?;
1032 this.update(&mut cx, |this, cx| {
1033 this.buffer
1034 .update(cx, |buffer, cx| buffer.set_language(Some(markdown), cx));
1035 })
1036 })
1037 .detach_and_log_err(cx);
1038 }
1039
1040 fn handle_buffer_event(
1041 &mut self,
1042 _: Model<Buffer>,
1043 event: &language::Event,
1044 cx: &mut ModelContext<Self>,
1045 ) {
1046 match event {
1047 language::Event::Operation(operation) => cx.emit(ContextEvent::Operation(
1048 ContextOperation::BufferOperation(operation.clone()),
1049 )),
1050 language::Event::Edited => {
1051 self.count_remaining_tokens(cx);
1052 self.reparse_slash_commands(cx);
1053 self.prune_invalid_edit_steps(cx);
1054 cx.emit(ContextEvent::MessagesEdited);
1055 }
1056 _ => {}
1057 }
1058 }
1059
1060 pub(crate) fn token_count(&self) -> Option<usize> {
1061 self.token_count
1062 }
1063
1064 pub(crate) fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
1065 let request = self.to_completion_request(cx);
1066 let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
1067 return;
1068 };
1069 self.pending_token_count = cx.spawn(|this, mut cx| {
1070 async move {
1071 cx.background_executor()
1072 .timer(Duration::from_millis(200))
1073 .await;
1074
1075 let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
1076 this.update(&mut cx, |this, cx| {
1077 this.token_count = Some(token_count);
1078 cx.notify()
1079 })
1080 }
1081 .log_err()
1082 });
1083 }
1084
1085 pub fn reparse_slash_commands(&mut self, cx: &mut ModelContext<Self>) {
1086 let buffer = self.buffer.read(cx);
1087 let mut row_ranges = self
1088 .edits_since_last_slash_command_parse
1089 .consume()
1090 .into_iter()
1091 .map(|edit| {
1092 let start_row = buffer.offset_to_point(edit.new.start).row;
1093 let end_row = buffer.offset_to_point(edit.new.end).row + 1;
1094 start_row..end_row
1095 })
1096 .peekable();
1097
1098 let mut removed = Vec::new();
1099 let mut updated = Vec::new();
1100 while let Some(mut row_range) = row_ranges.next() {
1101 while let Some(next_row_range) = row_ranges.peek() {
1102 if row_range.end >= next_row_range.start {
1103 row_range.end = next_row_range.end;
1104 row_ranges.next();
1105 } else {
1106 break;
1107 }
1108 }
1109
1110 let start = buffer.anchor_before(Point::new(row_range.start, 0));
1111 let end = buffer.anchor_after(Point::new(
1112 row_range.end - 1,
1113 buffer.line_len(row_range.end - 1),
1114 ));
1115
1116 let old_range = self.pending_command_indices_for_range(start..end, cx);
1117
1118 let mut new_commands = Vec::new();
1119 let mut lines = buffer.text_for_range(start..end).lines();
1120 let mut offset = lines.offset();
1121 while let Some(line) = lines.next() {
1122 if let Some(command_line) = SlashCommandLine::parse(line) {
1123 let name = &line[command_line.name.clone()];
1124 let argument = command_line.argument.as_ref().and_then(|argument| {
1125 (!argument.is_empty()).then_some(&line[argument.clone()])
1126 });
1127 if let Some(command) = SlashCommandRegistry::global(cx).command(name) {
1128 if !command.requires_argument() || argument.is_some() {
1129 let start_ix = offset + command_line.name.start - 1;
1130 let end_ix = offset
1131 + command_line
1132 .argument
1133 .map_or(command_line.name.end, |argument| argument.end);
1134 let source_range =
1135 buffer.anchor_after(start_ix)..buffer.anchor_after(end_ix);
1136 let pending_command = PendingSlashCommand {
1137 name: name.to_string(),
1138 argument: argument.map(ToString::to_string),
1139 source_range,
1140 status: PendingSlashCommandStatus::Idle,
1141 };
1142 updated.push(pending_command.clone());
1143 new_commands.push(pending_command);
1144 }
1145 }
1146 }
1147
1148 offset = lines.offset();
1149 }
1150
1151 let removed_commands = self.pending_slash_commands.splice(old_range, new_commands);
1152 removed.extend(removed_commands.map(|command| command.source_range));
1153 }
1154
1155 if !updated.is_empty() || !removed.is_empty() {
1156 cx.emit(ContextEvent::PendingSlashCommandsUpdated { removed, updated });
1157 }
1158 }
1159
1160 fn prune_invalid_edit_steps(&mut self, cx: &mut ModelContext<Self>) {
1161 let buffer = self.buffer.read(cx);
1162 let prev_len = self.workflow_steps.len();
1163 self.workflow_steps.retain(|step| {
1164 step.tagged_range.start.is_valid(buffer) && step.tagged_range.end.is_valid(buffer)
1165 });
1166 if self.workflow_steps.len() != prev_len {
1167 cx.emit(ContextEvent::WorkflowStepsChanged);
1168 cx.notify();
1169 }
1170 }
1171
1172 fn parse_edit_steps_in_range(
1173 &mut self,
1174 range: Range<usize>,
1175 project: Model<Project>,
1176 cx: &mut ModelContext<Self>,
1177 ) {
1178 let mut new_edit_steps = Vec::new();
1179
1180 let buffer = self.buffer.read(cx).snapshot();
1181 let mut message_lines = buffer.as_rope().chunks_in_range(range).lines();
1182 let mut in_step = false;
1183 let mut step_start = 0;
1184 let mut line_start_offset = message_lines.offset();
1185
1186 while let Some(line) = message_lines.next() {
1187 if let Some(step_start_index) = line.find("<step>") {
1188 if !in_step {
1189 in_step = true;
1190 step_start = line_start_offset + step_start_index;
1191 }
1192 }
1193
1194 if let Some(step_end_index) = line.find("</step>") {
1195 if in_step {
1196 let start_anchor = buffer.anchor_after(step_start);
1197 let end_anchor =
1198 buffer.anchor_before(line_start_offset + step_end_index + "</step>".len());
1199 let tagged_range = start_anchor..end_anchor;
1200
1201 // Check if a step with the same range already exists
1202 let existing_step_index = self
1203 .workflow_steps
1204 .binary_search_by(|probe| probe.tagged_range.cmp(&tagged_range, &buffer));
1205
1206 if let Err(ix) = existing_step_index {
1207 // Step doesn't exist, so add it
1208 let task = self.compute_workflow_step_edit_suggestions(
1209 tagged_range.clone(),
1210 project.clone(),
1211 cx,
1212 );
1213 new_edit_steps.push((
1214 ix,
1215 WorkflowStep {
1216 tagged_range,
1217 edit_suggestions: WorkflowStepEditSuggestions::Pending(task),
1218 },
1219 ));
1220 }
1221
1222 in_step = false;
1223 }
1224 }
1225
1226 line_start_offset = message_lines.offset();
1227 }
1228
1229 // Insert new steps and generate their corresponding tasks
1230 for (index, step) in new_edit_steps.into_iter().rev() {
1231 self.workflow_steps.insert(index, step);
1232 }
1233
1234 cx.emit(ContextEvent::WorkflowStepsChanged);
1235 cx.notify();
1236 }
1237
1238 fn compute_workflow_step_edit_suggestions(
1239 &self,
1240 tagged_range: Range<language::Anchor>,
1241 project: Model<Project>,
1242 cx: &mut ModelContext<Self>,
1243 ) -> Task<Option<()>> {
1244 let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
1245 return Task::ready(Err(anyhow!("no active model")).log_err());
1246 };
1247
1248 let mut request = self.to_completion_request(cx);
1249 let step_text = self
1250 .buffer
1251 .read(cx)
1252 .text_for_range(tagged_range.clone())
1253 .collect::<String>();
1254
1255 cx.spawn(|this, mut cx| {
1256 async move {
1257 let mut prompt = this.update(&mut cx, |this, _| {
1258 this.prompt_builder.generate_step_resolution_prompt()
1259 })??;
1260 prompt.push_str(&step_text);
1261
1262 request.messages.push(LanguageModelRequestMessage {
1263 role: Role::User,
1264 content: prompt,
1265 });
1266
1267 // Invoke the model to get its edit suggestions for this workflow step.
1268 let step_suggestions = model
1269 .use_tool::<tool::WorkflowStepEditSuggestions>(request, &cx)
1270 .await?;
1271
1272 // Translate the parsed suggestions to our internal types, which anchor the suggestions to locations in the code.
1273 let suggestion_tasks: Vec<_> = step_suggestions
1274 .edit_suggestions
1275 .iter()
1276 .map(|suggestion| suggestion.resolve(project.clone(), cx.clone()))
1277 .collect();
1278
1279 // Expand the context ranges of each suggestion and group suggestions with overlapping context ranges.
1280 let suggestions = future::join_all(suggestion_tasks)
1281 .await
1282 .into_iter()
1283 .filter_map(|task| task.log_err())
1284 .collect::<Vec<_>>();
1285
1286 let mut suggestions_by_buffer = HashMap::default();
1287 for (buffer, suggestion) in suggestions {
1288 suggestions_by_buffer
1289 .entry(buffer)
1290 .or_insert_with(Vec::new)
1291 .push(suggestion);
1292 }
1293
1294 let mut suggestion_groups_by_buffer = HashMap::default();
1295 for (buffer, mut suggestions) in suggestions_by_buffer {
1296 let mut suggestion_groups = Vec::<EditSuggestionGroup>::new();
1297 let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
1298 // Sort suggestions by their range so that earlier, larger ranges come first
1299 suggestions.sort_by(|a, b| a.range().cmp(&b.range(), &snapshot));
1300
1301 // Merge overlapping suggestions
1302 suggestions.dedup_by(|a, b| b.try_merge(&a, &snapshot));
1303
1304 // Create context ranges for each suggestion
1305 for suggestion in suggestions {
1306 let context_range = {
1307 let suggestion_point_range = suggestion.range().to_point(&snapshot);
1308 let start_row = suggestion_point_range.start.row.saturating_sub(5);
1309 let end_row = cmp::min(
1310 suggestion_point_range.end.row + 5,
1311 snapshot.max_point().row,
1312 );
1313 let start = snapshot.anchor_before(Point::new(start_row, 0));
1314 let end = snapshot
1315 .anchor_after(Point::new(end_row, snapshot.line_len(end_row)));
1316 start..end
1317 };
1318
1319 if let Some(last_group) = suggestion_groups.last_mut() {
1320 if last_group
1321 .context_range
1322 .end
1323 .cmp(&context_range.start, &snapshot)
1324 .is_ge()
1325 {
1326 // Merge with the previous group if context ranges overlap
1327 last_group.context_range.end = context_range.end;
1328 last_group.suggestions.push(suggestion);
1329 } else {
1330 // Create a new group
1331 suggestion_groups.push(EditSuggestionGroup {
1332 context_range,
1333 suggestions: vec![suggestion],
1334 });
1335 }
1336 } else {
1337 // Create the first group
1338 suggestion_groups.push(EditSuggestionGroup {
1339 context_range,
1340 suggestions: vec![suggestion],
1341 });
1342 }
1343 }
1344
1345 suggestion_groups_by_buffer.insert(buffer, suggestion_groups);
1346 }
1347
1348 this.update(&mut cx, |this, cx| {
1349 let step_index = this
1350 .workflow_steps
1351 .binary_search_by(|step| {
1352 step.tagged_range.cmp(&tagged_range, this.buffer.read(cx))
1353 })
1354 .map_err(|_| anyhow!("edit step not found"))?;
1355 if let Some(edit_step) = this.workflow_steps.get_mut(step_index) {
1356 edit_step.edit_suggestions = WorkflowStepEditSuggestions::Resolved(
1357 ResolvedWorkflowStepEditSuggestions {
1358 title: step_suggestions.step_title,
1359 edit_suggestions: suggestion_groups_by_buffer,
1360 },
1361 );
1362 cx.emit(ContextEvent::WorkflowStepsChanged);
1363 }
1364 anyhow::Ok(())
1365 })?
1366 }
1367 .log_err()
1368 })
1369 }
1370
1371 pub fn pending_command_for_position(
1372 &mut self,
1373 position: language::Anchor,
1374 cx: &mut ModelContext<Self>,
1375 ) -> Option<&mut PendingSlashCommand> {
1376 let buffer = self.buffer.read(cx);
1377 match self
1378 .pending_slash_commands
1379 .binary_search_by(|probe| probe.source_range.end.cmp(&position, buffer))
1380 {
1381 Ok(ix) => Some(&mut self.pending_slash_commands[ix]),
1382 Err(ix) => {
1383 let cmd = self.pending_slash_commands.get_mut(ix)?;
1384 if position.cmp(&cmd.source_range.start, buffer).is_ge()
1385 && position.cmp(&cmd.source_range.end, buffer).is_le()
1386 {
1387 Some(cmd)
1388 } else {
1389 None
1390 }
1391 }
1392 }
1393 }
1394
1395 pub fn pending_commands_for_range(
1396 &self,
1397 range: Range<language::Anchor>,
1398 cx: &AppContext,
1399 ) -> &[PendingSlashCommand] {
1400 let range = self.pending_command_indices_for_range(range, cx);
1401 &self.pending_slash_commands[range]
1402 }
1403
1404 fn pending_command_indices_for_range(
1405 &self,
1406 range: Range<language::Anchor>,
1407 cx: &AppContext,
1408 ) -> Range<usize> {
1409 let buffer = self.buffer.read(cx);
1410 let start_ix = match self
1411 .pending_slash_commands
1412 .binary_search_by(|probe| probe.source_range.end.cmp(&range.start, &buffer))
1413 {
1414 Ok(ix) | Err(ix) => ix,
1415 };
1416 let end_ix = match self
1417 .pending_slash_commands
1418 .binary_search_by(|probe| probe.source_range.start.cmp(&range.end, &buffer))
1419 {
1420 Ok(ix) => ix + 1,
1421 Err(ix) => ix,
1422 };
1423 start_ix..end_ix
1424 }
1425
1426 pub fn insert_command_output(
1427 &mut self,
1428 command_range: Range<language::Anchor>,
1429 output: Task<Result<SlashCommandOutput>>,
1430 insert_trailing_newline: bool,
1431 cx: &mut ModelContext<Self>,
1432 ) {
1433 self.reparse_slash_commands(cx);
1434
1435 let insert_output_task = cx.spawn(|this, mut cx| {
1436 let command_range = command_range.clone();
1437 async move {
1438 let output = output.await;
1439 this.update(&mut cx, |this, cx| match output {
1440 Ok(mut output) => {
1441 if insert_trailing_newline {
1442 output.text.push('\n');
1443 }
1444
1445 let version = this.version.clone();
1446 let command_id = SlashCommandId(this.next_timestamp());
1447 let (operation, event) = this.buffer.update(cx, |buffer, cx| {
1448 let start = command_range.start.to_offset(buffer);
1449 let old_end = command_range.end.to_offset(buffer);
1450 let new_end = start + output.text.len();
1451 buffer.edit([(start..old_end, output.text)], None, cx);
1452
1453 let mut sections = output
1454 .sections
1455 .into_iter()
1456 .map(|section| SlashCommandOutputSection {
1457 range: buffer.anchor_after(start + section.range.start)
1458 ..buffer.anchor_before(start + section.range.end),
1459 icon: section.icon,
1460 label: section.label,
1461 })
1462 .collect::<Vec<_>>();
1463 sections.sort_by(|a, b| a.range.cmp(&b.range, buffer));
1464
1465 this.slash_command_output_sections
1466 .extend(sections.iter().cloned());
1467 this.slash_command_output_sections
1468 .sort_by(|a, b| a.range.cmp(&b.range, buffer));
1469
1470 let output_range =
1471 buffer.anchor_after(start)..buffer.anchor_before(new_end);
1472 this.finished_slash_commands.insert(command_id);
1473
1474 (
1475 ContextOperation::SlashCommandFinished {
1476 id: command_id,
1477 output_range: output_range.clone(),
1478 sections: sections.clone(),
1479 version,
1480 },
1481 ContextEvent::SlashCommandFinished {
1482 output_range,
1483 sections,
1484 run_commands_in_output: output.run_commands_in_text,
1485 },
1486 )
1487 });
1488
1489 this.push_op(operation, cx);
1490 cx.emit(event);
1491 }
1492 Err(error) => {
1493 if let Some(pending_command) =
1494 this.pending_command_for_position(command_range.start, cx)
1495 {
1496 pending_command.status =
1497 PendingSlashCommandStatus::Error(error.to_string());
1498 cx.emit(ContextEvent::PendingSlashCommandsUpdated {
1499 removed: vec![pending_command.source_range.clone()],
1500 updated: vec![pending_command.clone()],
1501 });
1502 }
1503 }
1504 })
1505 .ok();
1506 }
1507 });
1508
1509 if let Some(pending_command) = self.pending_command_for_position(command_range.start, cx) {
1510 pending_command.status = PendingSlashCommandStatus::Running {
1511 _task: insert_output_task.shared(),
1512 };
1513 cx.emit(ContextEvent::PendingSlashCommandsUpdated {
1514 removed: vec![pending_command.source_range.clone()],
1515 updated: vec![pending_command.clone()],
1516 });
1517 }
1518 }
1519
1520 pub fn completion_provider_changed(&mut self, cx: &mut ModelContext<Self>) {
1521 self.count_remaining_tokens(cx);
1522 }
1523
1524 pub fn assist(&mut self, cx: &mut ModelContext<Self>) -> Option<MessageAnchor> {
1525 let provider = LanguageModelRegistry::read_global(cx).active_provider()?;
1526 let model = LanguageModelRegistry::read_global(cx).active_model()?;
1527 let last_message_id = self.message_anchors.iter().rev().find_map(|message| {
1528 message
1529 .start
1530 .is_valid(self.buffer.read(cx))
1531 .then_some(message.id)
1532 })?;
1533
1534 if !provider.is_authenticated(cx) {
1535 log::info!("completion provider has no credentials");
1536 return None;
1537 }
1538
1539 let request = self.to_completion_request(cx);
1540 let assistant_message = self
1541 .insert_message_after(last_message_id, Role::Assistant, MessageStatus::Pending, cx)
1542 .unwrap();
1543
1544 // Queue up the user's next reply.
1545 let user_message = self
1546 .insert_message_after(assistant_message.id, Role::User, MessageStatus::Done, cx)
1547 .unwrap();
1548
1549 let task = cx.spawn({
1550 |this, mut cx| async move {
1551 let stream = model.stream_completion(request, &cx);
1552 let assistant_message_id = assistant_message.id;
1553 let mut response_latency = None;
1554 let stream_completion = async {
1555 let request_start = Instant::now();
1556 let mut chunks = stream.await?;
1557
1558 while let Some(chunk) = chunks.next().await {
1559 if response_latency.is_none() {
1560 response_latency = Some(request_start.elapsed());
1561 }
1562 let chunk = chunk?;
1563
1564 this.update(&mut cx, |this, cx| {
1565 let message_ix = this
1566 .message_anchors
1567 .iter()
1568 .position(|message| message.id == assistant_message_id)?;
1569 let message_range = this.buffer.update(cx, |buffer, cx| {
1570 let message_start_offset =
1571 this.message_anchors[message_ix].start.to_offset(buffer);
1572 let message_old_end_offset = this.message_anchors[message_ix + 1..]
1573 .iter()
1574 .find(|message| message.start.is_valid(buffer))
1575 .map_or(buffer.len(), |message| {
1576 message.start.to_offset(buffer).saturating_sub(1)
1577 });
1578 let message_new_end_offset = message_old_end_offset + chunk.len();
1579 buffer.edit(
1580 [(message_old_end_offset..message_old_end_offset, chunk)],
1581 None,
1582 cx,
1583 );
1584 message_start_offset..message_new_end_offset
1585 });
1586 if let Some(project) = this.project.clone() {
1587 this.parse_edit_steps_in_range(message_range, project, cx);
1588 }
1589 cx.emit(ContextEvent::StreamedCompletion);
1590
1591 Some(())
1592 })?;
1593 smol::future::yield_now().await;
1594 }
1595
1596 this.update(&mut cx, |this, cx| {
1597 this.pending_completions
1598 .retain(|completion| completion.id != this.completion_count);
1599 this.summarize(false, cx);
1600 })?;
1601
1602 anyhow::Ok(())
1603 };
1604
1605 let result = stream_completion.await;
1606
1607 this.update(&mut cx, |this, cx| {
1608 let error_message = result
1609 .err()
1610 .map(|error| error.to_string().trim().to_string());
1611
1612 if let Some(error_message) = error_message.as_ref() {
1613 cx.emit(ContextEvent::AssistError(error_message.to_string()));
1614 }
1615
1616 this.update_metadata(assistant_message_id, cx, |metadata| {
1617 if let Some(error_message) = error_message.as_ref() {
1618 metadata.status =
1619 MessageStatus::Error(SharedString::from(error_message.clone()));
1620 } else {
1621 metadata.status = MessageStatus::Done;
1622 }
1623 });
1624
1625 if let Some(telemetry) = this.telemetry.as_ref() {
1626 telemetry.report_assistant_event(
1627 Some(this.id.0.clone()),
1628 AssistantKind::Panel,
1629 model.telemetry_id(),
1630 response_latency,
1631 error_message,
1632 );
1633 }
1634 })
1635 .ok();
1636 }
1637 });
1638
1639 self.pending_completions.push(PendingCompletion {
1640 id: post_inc(&mut self.completion_count),
1641 _task: task,
1642 });
1643
1644 Some(user_message)
1645 }
1646
1647 pub fn to_completion_request(&self, cx: &AppContext) -> LanguageModelRequest {
1648 let messages = self
1649 .messages(cx)
1650 .filter(|message| matches!(message.status, MessageStatus::Done))
1651 .map(|message| message.to_request_message(self.buffer.read(cx)));
1652
1653 LanguageModelRequest {
1654 messages: messages.collect(),
1655 stop: vec![],
1656 temperature: 1.0,
1657 }
1658 }
1659
1660 pub fn cancel_last_assist(&mut self) -> bool {
1661 self.pending_completions.pop().is_some()
1662 }
1663
1664 pub fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut ModelContext<Self>) {
1665 for id in ids {
1666 if let Some(metadata) = self.messages_metadata.get(&id) {
1667 let role = metadata.role.cycle();
1668 self.update_metadata(id, cx, |metadata| metadata.role = role);
1669 }
1670 }
1671 }
1672
1673 pub fn update_metadata(
1674 &mut self,
1675 id: MessageId,
1676 cx: &mut ModelContext<Self>,
1677 f: impl FnOnce(&mut MessageMetadata),
1678 ) {
1679 let version = self.version.clone();
1680 let timestamp = self.next_timestamp();
1681 if let Some(metadata) = self.messages_metadata.get_mut(&id) {
1682 f(metadata);
1683 metadata.timestamp = timestamp;
1684 let operation = ContextOperation::UpdateMessage {
1685 message_id: id,
1686 metadata: metadata.clone(),
1687 version,
1688 };
1689 self.push_op(operation, cx);
1690 cx.emit(ContextEvent::MessagesEdited);
1691 cx.notify();
1692 }
1693 }
1694
1695 fn insert_message_after(
1696 &mut self,
1697 message_id: MessageId,
1698 role: Role,
1699 status: MessageStatus,
1700 cx: &mut ModelContext<Self>,
1701 ) -> Option<MessageAnchor> {
1702 if let Some(prev_message_ix) = self
1703 .message_anchors
1704 .iter()
1705 .position(|message| message.id == message_id)
1706 {
1707 // Find the next valid message after the one we were given.
1708 let mut next_message_ix = prev_message_ix + 1;
1709 while let Some(next_message) = self.message_anchors.get(next_message_ix) {
1710 if next_message.start.is_valid(self.buffer.read(cx)) {
1711 break;
1712 }
1713 next_message_ix += 1;
1714 }
1715
1716 let start = self.buffer.update(cx, |buffer, cx| {
1717 let offset = self
1718 .message_anchors
1719 .get(next_message_ix)
1720 .map_or(buffer.len(), |message| {
1721 buffer.clip_offset(message.start.to_offset(buffer) - 1, Bias::Left)
1722 });
1723 buffer.edit([(offset..offset, "\n")], None, cx);
1724 buffer.anchor_before(offset + 1)
1725 });
1726
1727 let version = self.version.clone();
1728 let anchor = MessageAnchor {
1729 id: MessageId(self.next_timestamp()),
1730 start,
1731 };
1732 let metadata = MessageMetadata {
1733 role,
1734 status,
1735 timestamp: anchor.id.0,
1736 };
1737 self.insert_message(anchor.clone(), metadata.clone(), cx);
1738 self.push_op(
1739 ContextOperation::InsertMessage {
1740 anchor: anchor.clone(),
1741 metadata,
1742 version,
1743 },
1744 cx,
1745 );
1746 Some(anchor)
1747 } else {
1748 None
1749 }
1750 }
1751
1752 pub fn split_message(
1753 &mut self,
1754 range: Range<usize>,
1755 cx: &mut ModelContext<Self>,
1756 ) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
1757 let start_message = self.message_for_offset(range.start, cx);
1758 let end_message = self.message_for_offset(range.end, cx);
1759 if let Some((start_message, end_message)) = start_message.zip(end_message) {
1760 // Prevent splitting when range spans multiple messages.
1761 if start_message.id != end_message.id {
1762 return (None, None);
1763 }
1764
1765 let message = start_message;
1766 let role = message.role;
1767 let mut edited_buffer = false;
1768
1769 let mut suffix_start = None;
1770 if range.start > message.offset_range.start && range.end < message.offset_range.end - 1
1771 {
1772 if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
1773 suffix_start = Some(range.end + 1);
1774 } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
1775 suffix_start = Some(range.end);
1776 }
1777 }
1778
1779 let version = self.version.clone();
1780 let suffix = if let Some(suffix_start) = suffix_start {
1781 MessageAnchor {
1782 id: MessageId(self.next_timestamp()),
1783 start: self.buffer.read(cx).anchor_before(suffix_start),
1784 }
1785 } else {
1786 self.buffer.update(cx, |buffer, cx| {
1787 buffer.edit([(range.end..range.end, "\n")], None, cx);
1788 });
1789 edited_buffer = true;
1790 MessageAnchor {
1791 id: MessageId(self.next_timestamp()),
1792 start: self.buffer.read(cx).anchor_before(range.end + 1),
1793 }
1794 };
1795
1796 let suffix_metadata = MessageMetadata {
1797 role,
1798 status: MessageStatus::Done,
1799 timestamp: suffix.id.0,
1800 };
1801 self.insert_message(suffix.clone(), suffix_metadata.clone(), cx);
1802 self.push_op(
1803 ContextOperation::InsertMessage {
1804 anchor: suffix.clone(),
1805 metadata: suffix_metadata,
1806 version,
1807 },
1808 cx,
1809 );
1810
1811 let new_messages =
1812 if range.start == range.end || range.start == message.offset_range.start {
1813 (None, Some(suffix))
1814 } else {
1815 let mut prefix_end = None;
1816 if range.start > message.offset_range.start
1817 && range.end < message.offset_range.end - 1
1818 {
1819 if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
1820 prefix_end = Some(range.start + 1);
1821 } else if self.buffer.read(cx).reversed_chars_at(range.start).next()
1822 == Some('\n')
1823 {
1824 prefix_end = Some(range.start);
1825 }
1826 }
1827
1828 let version = self.version.clone();
1829 let selection = if let Some(prefix_end) = prefix_end {
1830 MessageAnchor {
1831 id: MessageId(self.next_timestamp()),
1832 start: self.buffer.read(cx).anchor_before(prefix_end),
1833 }
1834 } else {
1835 self.buffer.update(cx, |buffer, cx| {
1836 buffer.edit([(range.start..range.start, "\n")], None, cx)
1837 });
1838 edited_buffer = true;
1839 MessageAnchor {
1840 id: MessageId(self.next_timestamp()),
1841 start: self.buffer.read(cx).anchor_before(range.end + 1),
1842 }
1843 };
1844
1845 let selection_metadata = MessageMetadata {
1846 role,
1847 status: MessageStatus::Done,
1848 timestamp: selection.id.0,
1849 };
1850 self.insert_message(selection.clone(), selection_metadata.clone(), cx);
1851 self.push_op(
1852 ContextOperation::InsertMessage {
1853 anchor: selection.clone(),
1854 metadata: selection_metadata,
1855 version,
1856 },
1857 cx,
1858 );
1859
1860 (Some(selection), Some(suffix))
1861 };
1862
1863 if !edited_buffer {
1864 cx.emit(ContextEvent::MessagesEdited);
1865 }
1866 new_messages
1867 } else {
1868 (None, None)
1869 }
1870 }
1871
1872 fn insert_message(
1873 &mut self,
1874 new_anchor: MessageAnchor,
1875 new_metadata: MessageMetadata,
1876 cx: &mut ModelContext<Self>,
1877 ) {
1878 cx.emit(ContextEvent::MessagesEdited);
1879
1880 self.messages_metadata.insert(new_anchor.id, new_metadata);
1881
1882 let buffer = self.buffer.read(cx);
1883 let insertion_ix = self
1884 .message_anchors
1885 .iter()
1886 .position(|anchor| {
1887 let comparison = new_anchor.start.cmp(&anchor.start, buffer);
1888 comparison.is_lt() || (comparison.is_eq() && new_anchor.id > anchor.id)
1889 })
1890 .unwrap_or(self.message_anchors.len());
1891 self.message_anchors.insert(insertion_ix, new_anchor);
1892 }
1893
1894 pub(super) fn summarize(&mut self, replace_old: bool, cx: &mut ModelContext<Self>) {
1895 let Some(provider) = LanguageModelRegistry::read_global(cx).active_provider() else {
1896 return;
1897 };
1898 let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
1899 return;
1900 };
1901
1902 if replace_old || (self.message_anchors.len() >= 2 && self.summary.is_none()) {
1903 if !provider.is_authenticated(cx) {
1904 return;
1905 }
1906
1907 let messages = self
1908 .messages(cx)
1909 .map(|message| message.to_request_message(self.buffer.read(cx)))
1910 .chain(Some(LanguageModelRequestMessage {
1911 role: Role::User,
1912 content: "Summarize the context into a short title without punctuation.".into(),
1913 }));
1914 let request = LanguageModelRequest {
1915 messages: messages.collect(),
1916 stop: vec![],
1917 temperature: 1.0,
1918 };
1919
1920 self.pending_summary = cx.spawn(|this, mut cx| {
1921 async move {
1922 let stream = model.stream_completion(request, &cx);
1923 let mut messages = stream.await?;
1924
1925 let mut replaced = !replace_old;
1926 while let Some(message) = messages.next().await {
1927 let text = message?;
1928 let mut lines = text.lines();
1929 this.update(&mut cx, |this, cx| {
1930 let version = this.version.clone();
1931 let timestamp = this.next_timestamp();
1932 let summary = this.summary.get_or_insert(ContextSummary::default());
1933 if !replaced && replace_old {
1934 summary.text.clear();
1935 replaced = true;
1936 }
1937 summary.text.extend(lines.next());
1938 summary.timestamp = timestamp;
1939 let operation = ContextOperation::UpdateSummary {
1940 summary: summary.clone(),
1941 version,
1942 };
1943 this.push_op(operation, cx);
1944 cx.emit(ContextEvent::SummaryChanged);
1945 })?;
1946
1947 // Stop if the LLM generated multiple lines.
1948 if lines.next().is_some() {
1949 break;
1950 }
1951 }
1952
1953 this.update(&mut cx, |this, cx| {
1954 let version = this.version.clone();
1955 let timestamp = this.next_timestamp();
1956 if let Some(summary) = this.summary.as_mut() {
1957 summary.done = true;
1958 summary.timestamp = timestamp;
1959 let operation = ContextOperation::UpdateSummary {
1960 summary: summary.clone(),
1961 version,
1962 };
1963 this.push_op(operation, cx);
1964 cx.emit(ContextEvent::SummaryChanged);
1965 }
1966 })?;
1967
1968 anyhow::Ok(())
1969 }
1970 .log_err()
1971 });
1972 }
1973 }
1974
1975 fn message_for_offset(&self, offset: usize, cx: &AppContext) -> Option<Message> {
1976 self.messages_for_offsets([offset], cx).pop()
1977 }
1978
1979 pub fn messages_for_offsets(
1980 &self,
1981 offsets: impl IntoIterator<Item = usize>,
1982 cx: &AppContext,
1983 ) -> Vec<Message> {
1984 let mut result = Vec::new();
1985
1986 let mut messages = self.messages(cx).peekable();
1987 let mut offsets = offsets.into_iter().peekable();
1988 let mut current_message = messages.next();
1989 while let Some(offset) = offsets.next() {
1990 // Locate the message that contains the offset.
1991 while current_message.as_ref().map_or(false, |message| {
1992 !message.offset_range.contains(&offset) && messages.peek().is_some()
1993 }) {
1994 current_message = messages.next();
1995 }
1996 let Some(message) = current_message.as_ref() else {
1997 break;
1998 };
1999
2000 // Skip offsets that are in the same message.
2001 while offsets.peek().map_or(false, |offset| {
2002 message.offset_range.contains(offset) || messages.peek().is_none()
2003 }) {
2004 offsets.next();
2005 }
2006
2007 result.push(message.clone());
2008 }
2009 result
2010 }
2011
2012 pub fn messages<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator<Item = Message> {
2013 let buffer = self.buffer.read(cx);
2014 let mut message_anchors = self.message_anchors.iter().enumerate().peekable();
2015 iter::from_fn(move || {
2016 if let Some((start_ix, message_anchor)) = message_anchors.next() {
2017 let metadata = self.messages_metadata.get(&message_anchor.id)?;
2018 let message_start = message_anchor.start.to_offset(buffer);
2019 let mut message_end = None;
2020 let mut end_ix = start_ix;
2021 while let Some((_, next_message)) = message_anchors.peek() {
2022 if next_message.start.is_valid(buffer) {
2023 message_end = Some(next_message.start);
2024 break;
2025 } else {
2026 end_ix += 1;
2027 message_anchors.next();
2028 }
2029 }
2030 let message_end = message_end
2031 .unwrap_or(language::Anchor::MAX)
2032 .to_offset(buffer);
2033
2034 return Some(Message {
2035 index_range: start_ix..end_ix,
2036 offset_range: message_start..message_end,
2037 id: message_anchor.id,
2038 anchor: message_anchor.start,
2039 role: metadata.role,
2040 status: metadata.status.clone(),
2041 });
2042 }
2043 None
2044 })
2045 }
2046
2047 pub fn save(
2048 &mut self,
2049 debounce: Option<Duration>,
2050 fs: Arc<dyn Fs>,
2051 cx: &mut ModelContext<Context>,
2052 ) {
2053 if self.replica_id() != ReplicaId::default() {
2054 // Prevent saving a remote context for now.
2055 return;
2056 }
2057
2058 self.pending_save = cx.spawn(|this, mut cx| async move {
2059 if let Some(debounce) = debounce {
2060 cx.background_executor().timer(debounce).await;
2061 }
2062
2063 let (old_path, summary) = this.read_with(&cx, |this, _| {
2064 let path = this.path.clone();
2065 let summary = if let Some(summary) = this.summary.as_ref() {
2066 if summary.done {
2067 Some(summary.text.clone())
2068 } else {
2069 None
2070 }
2071 } else {
2072 None
2073 };
2074 (path, summary)
2075 })?;
2076
2077 if let Some(summary) = summary {
2078 let context = this.read_with(&cx, |this, cx| this.serialize(cx))?;
2079 let mut discriminant = 1;
2080 let mut new_path;
2081 loop {
2082 new_path = contexts_dir().join(&format!(
2083 "{} - {}.zed.json",
2084 summary.trim(),
2085 discriminant
2086 ));
2087 if fs.is_file(&new_path).await {
2088 discriminant += 1;
2089 } else {
2090 break;
2091 }
2092 }
2093
2094 fs.create_dir(contexts_dir().as_ref()).await?;
2095 fs.atomic_write(new_path.clone(), serde_json::to_string(&context).unwrap())
2096 .await?;
2097 if let Some(old_path) = old_path {
2098 if new_path != old_path {
2099 fs.remove_file(
2100 &old_path,
2101 RemoveOptions {
2102 recursive: false,
2103 ignore_if_not_exists: true,
2104 },
2105 )
2106 .await?;
2107 }
2108 }
2109
2110 this.update(&mut cx, |this, _| this.path = Some(new_path))?;
2111 }
2112
2113 Ok(())
2114 });
2115 }
2116
2117 pub(crate) fn custom_summary(&mut self, custom_summary: String, cx: &mut ModelContext<Self>) {
2118 let timestamp = self.next_timestamp();
2119 let summary = self.summary.get_or_insert(ContextSummary::default());
2120 summary.timestamp = timestamp;
2121 summary.done = true;
2122 summary.text = custom_summary;
2123 cx.emit(ContextEvent::SummaryChanged);
2124 }
2125}
2126
2127#[derive(Debug, Default)]
2128pub struct ContextVersion {
2129 context: clock::Global,
2130 buffer: clock::Global,
2131}
2132
2133impl ContextVersion {
2134 pub fn from_proto(proto: &proto::ContextVersion) -> Self {
2135 Self {
2136 context: language::proto::deserialize_version(&proto.context_version),
2137 buffer: language::proto::deserialize_version(&proto.buffer_version),
2138 }
2139 }
2140
2141 pub fn to_proto(&self, context_id: ContextId) -> proto::ContextVersion {
2142 proto::ContextVersion {
2143 context_id: context_id.to_proto(),
2144 context_version: language::proto::serialize_version(&self.context),
2145 buffer_version: language::proto::serialize_version(&self.buffer),
2146 }
2147 }
2148}
2149
2150#[derive(Debug, Clone)]
2151pub struct PendingSlashCommand {
2152 pub name: String,
2153 pub argument: Option<String>,
2154 pub status: PendingSlashCommandStatus,
2155 pub source_range: Range<language::Anchor>,
2156}
2157
2158#[derive(Debug, Clone)]
2159pub enum PendingSlashCommandStatus {
2160 Idle,
2161 Running { _task: Shared<Task<()>> },
2162 Error(String),
2163}
2164
2165#[derive(Serialize, Deserialize)]
2166pub struct SavedMessage {
2167 pub id: MessageId,
2168 pub start: usize,
2169 pub metadata: MessageMetadata,
2170}
2171
2172#[derive(Serialize, Deserialize)]
2173pub struct SavedContext {
2174 pub id: Option<ContextId>,
2175 pub zed: String,
2176 pub version: String,
2177 pub text: String,
2178 pub messages: Vec<SavedMessage>,
2179 pub summary: String,
2180 pub slash_command_output_sections:
2181 Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
2182}
2183
2184impl SavedContext {
2185 pub const VERSION: &'static str = "0.4.0";
2186
2187 pub fn from_json(json: &str) -> Result<Self> {
2188 let saved_context_json = serde_json::from_str::<serde_json::Value>(json)?;
2189 match saved_context_json
2190 .get("version")
2191 .ok_or_else(|| anyhow!("version not found"))?
2192 {
2193 serde_json::Value::String(version) => match version.as_str() {
2194 SavedContext::VERSION => {
2195 Ok(serde_json::from_value::<SavedContext>(saved_context_json)?)
2196 }
2197 SavedContextV0_3_0::VERSION => {
2198 let saved_context =
2199 serde_json::from_value::<SavedContextV0_3_0>(saved_context_json)?;
2200 Ok(saved_context.upgrade())
2201 }
2202 SavedContextV0_2_0::VERSION => {
2203 let saved_context =
2204 serde_json::from_value::<SavedContextV0_2_0>(saved_context_json)?;
2205 Ok(saved_context.upgrade())
2206 }
2207 SavedContextV0_1_0::VERSION => {
2208 let saved_context =
2209 serde_json::from_value::<SavedContextV0_1_0>(saved_context_json)?;
2210 Ok(saved_context.upgrade())
2211 }
2212 _ => Err(anyhow!("unrecognized saved context version: {}", version)),
2213 },
2214 _ => Err(anyhow!("version not found on saved context")),
2215 }
2216 }
2217
2218 fn into_ops(
2219 self,
2220 buffer: &Model<Buffer>,
2221 cx: &mut ModelContext<Context>,
2222 ) -> Vec<ContextOperation> {
2223 let mut operations = Vec::new();
2224 let mut version = clock::Global::new();
2225 let mut next_timestamp = clock::Lamport::new(ReplicaId::default());
2226
2227 let mut first_message_metadata = None;
2228 for message in self.messages {
2229 if message.id == MessageId(clock::Lamport::default()) {
2230 first_message_metadata = Some(message.metadata);
2231 } else {
2232 operations.push(ContextOperation::InsertMessage {
2233 anchor: MessageAnchor {
2234 id: message.id,
2235 start: buffer.read(cx).anchor_before(message.start),
2236 },
2237 metadata: MessageMetadata {
2238 role: message.metadata.role,
2239 status: message.metadata.status,
2240 timestamp: message.metadata.timestamp,
2241 },
2242 version: version.clone(),
2243 });
2244 version.observe(message.id.0);
2245 next_timestamp.observe(message.id.0);
2246 }
2247 }
2248
2249 if let Some(metadata) = first_message_metadata {
2250 let timestamp = next_timestamp.tick();
2251 operations.push(ContextOperation::UpdateMessage {
2252 message_id: MessageId(clock::Lamport::default()),
2253 metadata: MessageMetadata {
2254 role: metadata.role,
2255 status: metadata.status,
2256 timestamp,
2257 },
2258 version: version.clone(),
2259 });
2260 version.observe(timestamp);
2261 }
2262
2263 let timestamp = next_timestamp.tick();
2264 operations.push(ContextOperation::SlashCommandFinished {
2265 id: SlashCommandId(timestamp),
2266 output_range: language::Anchor::MIN..language::Anchor::MAX,
2267 sections: self
2268 .slash_command_output_sections
2269 .into_iter()
2270 .map(|section| {
2271 let buffer = buffer.read(cx);
2272 SlashCommandOutputSection {
2273 range: buffer.anchor_after(section.range.start)
2274 ..buffer.anchor_before(section.range.end),
2275 icon: section.icon,
2276 label: section.label,
2277 }
2278 })
2279 .collect(),
2280 version: version.clone(),
2281 });
2282 version.observe(timestamp);
2283
2284 let timestamp = next_timestamp.tick();
2285 operations.push(ContextOperation::UpdateSummary {
2286 summary: ContextSummary {
2287 text: self.summary,
2288 done: true,
2289 timestamp,
2290 },
2291 version: version.clone(),
2292 });
2293 version.observe(timestamp);
2294
2295 operations
2296 }
2297}
2298
2299#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
2300struct SavedMessageIdPreV0_4_0(usize);
2301
2302#[derive(Serialize, Deserialize)]
2303struct SavedMessagePreV0_4_0 {
2304 id: SavedMessageIdPreV0_4_0,
2305 start: usize,
2306}
2307
2308#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
2309struct SavedMessageMetadataPreV0_4_0 {
2310 role: Role,
2311 status: MessageStatus,
2312}
2313
2314#[derive(Serialize, Deserialize)]
2315struct SavedContextV0_3_0 {
2316 id: Option<ContextId>,
2317 zed: String,
2318 version: String,
2319 text: String,
2320 messages: Vec<SavedMessagePreV0_4_0>,
2321 message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
2322 summary: String,
2323 slash_command_output_sections: Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
2324}
2325
2326impl SavedContextV0_3_0 {
2327 const VERSION: &'static str = "0.3.0";
2328
2329 fn upgrade(self) -> SavedContext {
2330 SavedContext {
2331 id: self.id,
2332 zed: self.zed,
2333 version: SavedContext::VERSION.into(),
2334 text: self.text,
2335 messages: self
2336 .messages
2337 .into_iter()
2338 .filter_map(|message| {
2339 let metadata = self.message_metadata.get(&message.id)?;
2340 let timestamp = clock::Lamport {
2341 replica_id: ReplicaId::default(),
2342 value: message.id.0 as u32,
2343 };
2344 Some(SavedMessage {
2345 id: MessageId(timestamp),
2346 start: message.start,
2347 metadata: MessageMetadata {
2348 role: metadata.role,
2349 status: metadata.status.clone(),
2350 timestamp,
2351 },
2352 })
2353 })
2354 .collect(),
2355 summary: self.summary,
2356 slash_command_output_sections: self.slash_command_output_sections,
2357 }
2358 }
2359}
2360
2361#[derive(Serialize, Deserialize)]
2362struct SavedContextV0_2_0 {
2363 id: Option<ContextId>,
2364 zed: String,
2365 version: String,
2366 text: String,
2367 messages: Vec<SavedMessagePreV0_4_0>,
2368 message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
2369 summary: String,
2370}
2371
2372impl SavedContextV0_2_0 {
2373 const VERSION: &'static str = "0.2.0";
2374
2375 fn upgrade(self) -> SavedContext {
2376 SavedContextV0_3_0 {
2377 id: self.id,
2378 zed: self.zed,
2379 version: SavedContextV0_3_0::VERSION.to_string(),
2380 text: self.text,
2381 messages: self.messages,
2382 message_metadata: self.message_metadata,
2383 summary: self.summary,
2384 slash_command_output_sections: Vec::new(),
2385 }
2386 .upgrade()
2387 }
2388}
2389
2390#[derive(Serialize, Deserialize)]
2391struct SavedContextV0_1_0 {
2392 id: Option<ContextId>,
2393 zed: String,
2394 version: String,
2395 text: String,
2396 messages: Vec<SavedMessagePreV0_4_0>,
2397 message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
2398 summary: String,
2399 api_url: Option<String>,
2400 model: OpenAiModel,
2401}
2402
2403impl SavedContextV0_1_0 {
2404 const VERSION: &'static str = "0.1.0";
2405
2406 fn upgrade(self) -> SavedContext {
2407 SavedContextV0_2_0 {
2408 id: self.id,
2409 zed: self.zed,
2410 version: SavedContextV0_2_0::VERSION.to_string(),
2411 text: self.text,
2412 messages: self.messages,
2413 message_metadata: self.message_metadata,
2414 summary: self.summary,
2415 }
2416 .upgrade()
2417 }
2418}
2419
2420#[derive(Clone)]
2421pub struct SavedContextMetadata {
2422 pub title: String,
2423 pub path: PathBuf,
2424 pub mtime: chrono::DateTime<chrono::Local>,
2425}
2426
2427#[cfg(test)]
2428mod tests {
2429 use super::*;
2430 use crate::{
2431 assistant_panel, prompt_library,
2432 slash_command::{active_command, file_command},
2433 MessageId,
2434 };
2435 use assistant_slash_command::{ArgumentCompletion, SlashCommand};
2436 use fs::FakeFs;
2437 use gpui::{AppContext, TestAppContext, WeakView};
2438 use indoc::indoc;
2439 use language::LspAdapterDelegate;
2440 use parking_lot::Mutex;
2441 use project::Project;
2442 use rand::prelude::*;
2443 use serde_json::json;
2444 use settings::SettingsStore;
2445 use std::{cell::RefCell, env, rc::Rc, sync::atomic::AtomicBool};
2446 use text::{network::Network, ToPoint};
2447 use ui::WindowContext;
2448 use unindent::Unindent;
2449 use util::{test::marked_text_ranges, RandomCharIter};
2450 use workspace::Workspace;
2451
2452 #[gpui::test]
2453 fn test_inserting_and_removing_messages(cx: &mut AppContext) {
2454 let settings_store = SettingsStore::test(cx);
2455 LanguageModelRegistry::test(cx);
2456 cx.set_global(settings_store);
2457 assistant_panel::init(cx);
2458 let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));
2459 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2460 let context =
2461 cx.new_model(|cx| Context::local(registry, None, None, prompt_builder.clone(), cx));
2462 let buffer = context.read(cx).buffer.clone();
2463
2464 let message_1 = context.read(cx).message_anchors[0].clone();
2465 assert_eq!(
2466 messages(&context, cx),
2467 vec![(message_1.id, Role::User, 0..0)]
2468 );
2469
2470 let message_2 = context.update(cx, |context, cx| {
2471 context
2472 .insert_message_after(message_1.id, Role::Assistant, MessageStatus::Done, cx)
2473 .unwrap()
2474 });
2475 assert_eq!(
2476 messages(&context, cx),
2477 vec![
2478 (message_1.id, Role::User, 0..1),
2479 (message_2.id, Role::Assistant, 1..1)
2480 ]
2481 );
2482
2483 buffer.update(cx, |buffer, cx| {
2484 buffer.edit([(0..0, "1"), (1..1, "2")], None, cx)
2485 });
2486 assert_eq!(
2487 messages(&context, cx),
2488 vec![
2489 (message_1.id, Role::User, 0..2),
2490 (message_2.id, Role::Assistant, 2..3)
2491 ]
2492 );
2493
2494 let message_3 = context.update(cx, |context, cx| {
2495 context
2496 .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
2497 .unwrap()
2498 });
2499 assert_eq!(
2500 messages(&context, cx),
2501 vec![
2502 (message_1.id, Role::User, 0..2),
2503 (message_2.id, Role::Assistant, 2..4),
2504 (message_3.id, Role::User, 4..4)
2505 ]
2506 );
2507
2508 let message_4 = context.update(cx, |context, cx| {
2509 context
2510 .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
2511 .unwrap()
2512 });
2513 assert_eq!(
2514 messages(&context, cx),
2515 vec![
2516 (message_1.id, Role::User, 0..2),
2517 (message_2.id, Role::Assistant, 2..4),
2518 (message_4.id, Role::User, 4..5),
2519 (message_3.id, Role::User, 5..5),
2520 ]
2521 );
2522
2523 buffer.update(cx, |buffer, cx| {
2524 buffer.edit([(4..4, "C"), (5..5, "D")], None, cx)
2525 });
2526 assert_eq!(
2527 messages(&context, cx),
2528 vec![
2529 (message_1.id, Role::User, 0..2),
2530 (message_2.id, Role::Assistant, 2..4),
2531 (message_4.id, Role::User, 4..6),
2532 (message_3.id, Role::User, 6..7),
2533 ]
2534 );
2535
2536 // Deleting across message boundaries merges the messages.
2537 buffer.update(cx, |buffer, cx| buffer.edit([(1..4, "")], None, cx));
2538 assert_eq!(
2539 messages(&context, cx),
2540 vec![
2541 (message_1.id, Role::User, 0..3),
2542 (message_3.id, Role::User, 3..4),
2543 ]
2544 );
2545
2546 // Undoing the deletion should also undo the merge.
2547 buffer.update(cx, |buffer, cx| buffer.undo(cx));
2548 assert_eq!(
2549 messages(&context, cx),
2550 vec![
2551 (message_1.id, Role::User, 0..2),
2552 (message_2.id, Role::Assistant, 2..4),
2553 (message_4.id, Role::User, 4..6),
2554 (message_3.id, Role::User, 6..7),
2555 ]
2556 );
2557
2558 // Redoing the deletion should also redo the merge.
2559 buffer.update(cx, |buffer, cx| buffer.redo(cx));
2560 assert_eq!(
2561 messages(&context, cx),
2562 vec![
2563 (message_1.id, Role::User, 0..3),
2564 (message_3.id, Role::User, 3..4),
2565 ]
2566 );
2567
2568 // Ensure we can still insert after a merged message.
2569 let message_5 = context.update(cx, |context, cx| {
2570 context
2571 .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)
2572 .unwrap()
2573 });
2574 assert_eq!(
2575 messages(&context, cx),
2576 vec![
2577 (message_1.id, Role::User, 0..3),
2578 (message_5.id, Role::System, 3..4),
2579 (message_3.id, Role::User, 4..5)
2580 ]
2581 );
2582 }
2583
2584 #[gpui::test]
2585 fn test_message_splitting(cx: &mut AppContext) {
2586 let settings_store = SettingsStore::test(cx);
2587 cx.set_global(settings_store);
2588 LanguageModelRegistry::test(cx);
2589 assistant_panel::init(cx);
2590 let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));
2591
2592 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2593 let context =
2594 cx.new_model(|cx| Context::local(registry, None, None, prompt_builder.clone(), cx));
2595 let buffer = context.read(cx).buffer.clone();
2596
2597 let message_1 = context.read(cx).message_anchors[0].clone();
2598 assert_eq!(
2599 messages(&context, cx),
2600 vec![(message_1.id, Role::User, 0..0)]
2601 );
2602
2603 buffer.update(cx, |buffer, cx| {
2604 buffer.edit([(0..0, "aaa\nbbb\nccc\nddd\n")], None, cx)
2605 });
2606
2607 let (_, message_2) = context.update(cx, |context, cx| context.split_message(3..3, cx));
2608 let message_2 = message_2.unwrap();
2609
2610 // We recycle newlines in the middle of a split message
2611 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\nddd\n");
2612 assert_eq!(
2613 messages(&context, cx),
2614 vec![
2615 (message_1.id, Role::User, 0..4),
2616 (message_2.id, Role::User, 4..16),
2617 ]
2618 );
2619
2620 let (_, message_3) = context.update(cx, |context, cx| context.split_message(3..3, cx));
2621 let message_3 = message_3.unwrap();
2622
2623 // We don't recycle newlines at the end of a split message
2624 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
2625 assert_eq!(
2626 messages(&context, cx),
2627 vec![
2628 (message_1.id, Role::User, 0..4),
2629 (message_3.id, Role::User, 4..5),
2630 (message_2.id, Role::User, 5..17),
2631 ]
2632 );
2633
2634 let (_, message_4) = context.update(cx, |context, cx| context.split_message(9..9, cx));
2635 let message_4 = message_4.unwrap();
2636 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
2637 assert_eq!(
2638 messages(&context, cx),
2639 vec![
2640 (message_1.id, Role::User, 0..4),
2641 (message_3.id, Role::User, 4..5),
2642 (message_2.id, Role::User, 5..9),
2643 (message_4.id, Role::User, 9..17),
2644 ]
2645 );
2646
2647 let (_, message_5) = context.update(cx, |context, cx| context.split_message(9..9, cx));
2648 let message_5 = message_5.unwrap();
2649 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\nddd\n");
2650 assert_eq!(
2651 messages(&context, cx),
2652 vec![
2653 (message_1.id, Role::User, 0..4),
2654 (message_3.id, Role::User, 4..5),
2655 (message_2.id, Role::User, 5..9),
2656 (message_4.id, Role::User, 9..10),
2657 (message_5.id, Role::User, 10..18),
2658 ]
2659 );
2660
2661 let (message_6, message_7) =
2662 context.update(cx, |context, cx| context.split_message(14..16, cx));
2663 let message_6 = message_6.unwrap();
2664 let message_7 = message_7.unwrap();
2665 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\ndd\nd\n");
2666 assert_eq!(
2667 messages(&context, cx),
2668 vec![
2669 (message_1.id, Role::User, 0..4),
2670 (message_3.id, Role::User, 4..5),
2671 (message_2.id, Role::User, 5..9),
2672 (message_4.id, Role::User, 9..10),
2673 (message_5.id, Role::User, 10..14),
2674 (message_6.id, Role::User, 14..17),
2675 (message_7.id, Role::User, 17..19),
2676 ]
2677 );
2678 }
2679
2680 #[gpui::test]
2681 fn test_messages_for_offsets(cx: &mut AppContext) {
2682 let settings_store = SettingsStore::test(cx);
2683 LanguageModelRegistry::test(cx);
2684 cx.set_global(settings_store);
2685 assistant_panel::init(cx);
2686 let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));
2687 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2688 let context =
2689 cx.new_model(|cx| Context::local(registry, None, None, prompt_builder.clone(), cx));
2690 let buffer = context.read(cx).buffer.clone();
2691
2692 let message_1 = context.read(cx).message_anchors[0].clone();
2693 assert_eq!(
2694 messages(&context, cx),
2695 vec![(message_1.id, Role::User, 0..0)]
2696 );
2697
2698 buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx));
2699 let message_2 = context
2700 .update(cx, |context, cx| {
2701 context.insert_message_after(message_1.id, Role::User, MessageStatus::Done, cx)
2702 })
2703 .unwrap();
2704 buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "bbb")], None, cx));
2705
2706 let message_3 = context
2707 .update(cx, |context, cx| {
2708 context.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
2709 })
2710 .unwrap();
2711 buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "ccc")], None, cx));
2712
2713 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc");
2714 assert_eq!(
2715 messages(&context, cx),
2716 vec![
2717 (message_1.id, Role::User, 0..4),
2718 (message_2.id, Role::User, 4..8),
2719 (message_3.id, Role::User, 8..11)
2720 ]
2721 );
2722
2723 assert_eq!(
2724 message_ids_for_offsets(&context, &[0, 4, 9], cx),
2725 [message_1.id, message_2.id, message_3.id]
2726 );
2727 assert_eq!(
2728 message_ids_for_offsets(&context, &[0, 1, 11], cx),
2729 [message_1.id, message_3.id]
2730 );
2731
2732 let message_4 = context
2733 .update(cx, |context, cx| {
2734 context.insert_message_after(message_3.id, Role::User, MessageStatus::Done, cx)
2735 })
2736 .unwrap();
2737 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\n");
2738 assert_eq!(
2739 messages(&context, cx),
2740 vec![
2741 (message_1.id, Role::User, 0..4),
2742 (message_2.id, Role::User, 4..8),
2743 (message_3.id, Role::User, 8..12),
2744 (message_4.id, Role::User, 12..12)
2745 ]
2746 );
2747 assert_eq!(
2748 message_ids_for_offsets(&context, &[0, 4, 8, 12], cx),
2749 [message_1.id, message_2.id, message_3.id, message_4.id]
2750 );
2751
2752 fn message_ids_for_offsets(
2753 context: &Model<Context>,
2754 offsets: &[usize],
2755 cx: &AppContext,
2756 ) -> Vec<MessageId> {
2757 context
2758 .read(cx)
2759 .messages_for_offsets(offsets.iter().copied(), cx)
2760 .into_iter()
2761 .map(|message| message.id)
2762 .collect()
2763 }
2764 }
2765
2766 #[gpui::test]
2767 async fn test_slash_commands(cx: &mut TestAppContext) {
2768 let settings_store = cx.update(SettingsStore::test);
2769 cx.set_global(settings_store);
2770 cx.update(LanguageModelRegistry::test);
2771 cx.update(Project::init_settings);
2772 cx.update(assistant_panel::init);
2773 let fs = FakeFs::new(cx.background_executor.clone());
2774
2775 fs.insert_tree(
2776 "/test",
2777 json!({
2778 "src": {
2779 "lib.rs": "fn one() -> usize { 1 }",
2780 "main.rs": "
2781 use crate::one;
2782 fn main() { one(); }
2783 ".unindent(),
2784 }
2785 }),
2786 )
2787 .await;
2788
2789 let slash_command_registry = cx.update(SlashCommandRegistry::default_global);
2790 slash_command_registry.register_command(file_command::FileSlashCommand, false);
2791 slash_command_registry.register_command(active_command::ActiveSlashCommand, false);
2792
2793 let registry = Arc::new(LanguageRegistry::test(cx.executor()));
2794 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2795 let context = cx.new_model(|cx| {
2796 Context::local(registry.clone(), None, None, prompt_builder.clone(), cx)
2797 });
2798
2799 let output_ranges = Rc::new(RefCell::new(HashSet::default()));
2800 context.update(cx, |_, cx| {
2801 cx.subscribe(&context, {
2802 let ranges = output_ranges.clone();
2803 move |_, _, event, _| match event {
2804 ContextEvent::PendingSlashCommandsUpdated { removed, updated } => {
2805 for range in removed {
2806 ranges.borrow_mut().remove(range);
2807 }
2808 for command in updated {
2809 ranges.borrow_mut().insert(command.source_range.clone());
2810 }
2811 }
2812 _ => {}
2813 }
2814 })
2815 .detach();
2816 });
2817
2818 let buffer = context.read_with(cx, |context, _| context.buffer.clone());
2819
2820 // Insert a slash command
2821 buffer.update(cx, |buffer, cx| {
2822 buffer.edit([(0..0, "/file src/lib.rs")], None, cx);
2823 });
2824 assert_text_and_output_ranges(
2825 &buffer,
2826 &output_ranges.borrow(),
2827 "
2828 «/file src/lib.rs»
2829 "
2830 .unindent()
2831 .trim_end(),
2832 cx,
2833 );
2834
2835 // Edit the argument of the slash command.
2836 buffer.update(cx, |buffer, cx| {
2837 let edit_offset = buffer.text().find("lib.rs").unwrap();
2838 buffer.edit([(edit_offset..edit_offset + "lib".len(), "main")], None, cx);
2839 });
2840 assert_text_and_output_ranges(
2841 &buffer,
2842 &output_ranges.borrow(),
2843 "
2844 «/file src/main.rs»
2845 "
2846 .unindent()
2847 .trim_end(),
2848 cx,
2849 );
2850
2851 // Edit the name of the slash command, using one that doesn't exist.
2852 buffer.update(cx, |buffer, cx| {
2853 let edit_offset = buffer.text().find("/file").unwrap();
2854 buffer.edit(
2855 [(edit_offset..edit_offset + "/file".len(), "/unknown")],
2856 None,
2857 cx,
2858 );
2859 });
2860 assert_text_and_output_ranges(
2861 &buffer,
2862 &output_ranges.borrow(),
2863 "
2864 /unknown src/main.rs
2865 "
2866 .unindent()
2867 .trim_end(),
2868 cx,
2869 );
2870
2871 #[track_caller]
2872 fn assert_text_and_output_ranges(
2873 buffer: &Model<Buffer>,
2874 ranges: &HashSet<Range<language::Anchor>>,
2875 expected_marked_text: &str,
2876 cx: &mut TestAppContext,
2877 ) {
2878 let (expected_text, expected_ranges) = marked_text_ranges(expected_marked_text, false);
2879 let (actual_text, actual_ranges) = buffer.update(cx, |buffer, _| {
2880 let mut ranges = ranges
2881 .iter()
2882 .map(|range| range.to_offset(buffer))
2883 .collect::<Vec<_>>();
2884 ranges.sort_by_key(|a| a.start);
2885 (buffer.text(), ranges)
2886 });
2887
2888 assert_eq!(actual_text, expected_text);
2889 assert_eq!(actual_ranges, expected_ranges);
2890 }
2891 }
2892
2893 #[gpui::test]
2894 async fn test_edit_step_parsing(cx: &mut TestAppContext) {
2895 cx.update(prompt_library::init);
2896 let settings_store = cx.update(SettingsStore::test);
2897 cx.set_global(settings_store);
2898 cx.update(Project::init_settings);
2899 let fs = FakeFs::new(cx.executor());
2900 fs.as_fake()
2901 .insert_tree(
2902 "/root",
2903 json!({
2904 "hello.rs": r#"
2905 fn hello() {
2906 println!("Hello, World!");
2907 }
2908 "#.unindent()
2909 }),
2910 )
2911 .await;
2912 let project = Project::test(fs, [Path::new("/root")], cx).await;
2913 cx.update(LanguageModelRegistry::test);
2914
2915 let model = cx.read(|cx| {
2916 LanguageModelRegistry::read_global(cx)
2917 .active_model()
2918 .unwrap()
2919 });
2920 cx.update(assistant_panel::init);
2921 let registry = Arc::new(LanguageRegistry::test(cx.executor()));
2922
2923 // Create a new context
2924 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2925 let context = cx.new_model(|cx| {
2926 Context::local(
2927 registry.clone(),
2928 Some(project),
2929 None,
2930 prompt_builder.clone(),
2931 cx,
2932 )
2933 });
2934 let buffer = context.read_with(cx, |context, _| context.buffer.clone());
2935
2936 // Simulate user input
2937 let user_message = indoc! {r#"
2938 Please add unnecessary complexity to this code:
2939
2940 ```hello.rs
2941 fn main() {
2942 println!("Hello, World!");
2943 }
2944 ```
2945 "#};
2946 buffer.update(cx, |buffer, cx| {
2947 buffer.edit([(0..0, user_message)], None, cx);
2948 });
2949
2950 // Simulate LLM response with edit steps
2951 let llm_response = indoc! {r#"
2952 Sure, I can help you with that. Here's a step-by-step process:
2953
2954 <step>
2955 First, let's extract the greeting into a separate function:
2956
2957 ```rust
2958 fn greet() {
2959 println!("Hello, World!");
2960 }
2961
2962 fn main() {
2963 greet();
2964 }
2965 ```
2966 </step>
2967
2968 <step>
2969 Now, let's make the greeting customizable:
2970
2971 ```rust
2972 fn greet(name: &str) {
2973 println!("Hello, {}!", name);
2974 }
2975
2976 fn main() {
2977 greet("World");
2978 }
2979 ```
2980 </step>
2981
2982 These changes make the code more modular and flexible.
2983 "#};
2984
2985 // Simulate the assist method to trigger the LLM response
2986 context.update(cx, |context, cx| context.assist(cx));
2987 cx.run_until_parked();
2988
2989 // Retrieve the assistant response message's start from the context
2990 let response_start_row = context.read_with(cx, |context, cx| {
2991 let buffer = context.buffer.read(cx);
2992 context.message_anchors[1].start.to_point(buffer).row
2993 });
2994
2995 // Simulate the LLM completion
2996 model
2997 .as_fake()
2998 .stream_last_completion_response(llm_response.to_string());
2999 model.as_fake().end_last_completion_stream();
3000
3001 // Wait for the completion to be processed
3002 cx.run_until_parked();
3003
3004 // Verify that the edit steps were parsed correctly
3005 context.read_with(cx, |context, cx| {
3006 assert_eq!(
3007 workflow_steps(context, cx),
3008 vec![
3009 (
3010 Point::new(response_start_row + 2, 0)
3011 ..Point::new(response_start_row + 14, 7),
3012 WorkflowStepEditSuggestionStatus::Pending
3013 ),
3014 (
3015 Point::new(response_start_row + 16, 0)
3016 ..Point::new(response_start_row + 28, 7),
3017 WorkflowStepEditSuggestionStatus::Pending
3018 ),
3019 ]
3020 );
3021 });
3022
3023 model
3024 .as_fake()
3025 .respond_to_last_tool_use(Ok(serde_json::to_value(
3026 tool::WorkflowStepEditSuggestions {
3027 step_title: "Title".into(),
3028 edit_suggestions: vec![tool::EditSuggestion {
3029 path: "/root/hello.rs".into(),
3030 // Simulate a symbol name that's slightly different than our outline query
3031 kind: tool::EditSuggestionKind::Update {
3032 symbol: "fn main()".into(),
3033 description: "Extract a greeting function".into(),
3034 },
3035 }],
3036 },
3037 )
3038 .unwrap()));
3039
3040 // Wait for tool use to be processed.
3041 cx.run_until_parked();
3042
3043 // Verify that the last edit step is not pending anymore.
3044 context.read_with(cx, |context, cx| {
3045 assert_eq!(
3046 workflow_steps(context, cx),
3047 vec![
3048 (
3049 Point::new(response_start_row + 2, 0)
3050 ..Point::new(response_start_row + 14, 7),
3051 WorkflowStepEditSuggestionStatus::Pending
3052 ),
3053 (
3054 Point::new(response_start_row + 16, 0)
3055 ..Point::new(response_start_row + 28, 7),
3056 WorkflowStepEditSuggestionStatus::Resolved
3057 ),
3058 ]
3059 );
3060 });
3061
3062 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
3063 enum WorkflowStepEditSuggestionStatus {
3064 Pending,
3065 Resolved,
3066 }
3067
3068 fn workflow_steps(
3069 context: &Context,
3070 cx: &AppContext,
3071 ) -> Vec<(Range<Point>, WorkflowStepEditSuggestionStatus)> {
3072 context
3073 .workflow_steps
3074 .iter()
3075 .map(|step| {
3076 let buffer = context.buffer.read(cx);
3077 let status = match &step.edit_suggestions {
3078 WorkflowStepEditSuggestions::Pending(_) => {
3079 WorkflowStepEditSuggestionStatus::Pending
3080 }
3081 WorkflowStepEditSuggestions::Resolved { .. } => {
3082 WorkflowStepEditSuggestionStatus::Resolved
3083 }
3084 };
3085 (step.tagged_range.to_point(buffer), status)
3086 })
3087 .collect()
3088 }
3089 }
3090
3091 #[gpui::test]
3092 async fn test_serialization(cx: &mut TestAppContext) {
3093 let settings_store = cx.update(SettingsStore::test);
3094 cx.set_global(settings_store);
3095 cx.update(LanguageModelRegistry::test);
3096 cx.update(assistant_panel::init);
3097 let registry = Arc::new(LanguageRegistry::test(cx.executor()));
3098 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3099 let context = cx.new_model(|cx| {
3100 Context::local(registry.clone(), None, None, prompt_builder.clone(), cx)
3101 });
3102 let buffer = context.read_with(cx, |context, _| context.buffer.clone());
3103 let message_0 = context.read_with(cx, |context, _| context.message_anchors[0].id);
3104 let message_1 = context.update(cx, |context, cx| {
3105 context
3106 .insert_message_after(message_0, Role::Assistant, MessageStatus::Done, cx)
3107 .unwrap()
3108 });
3109 let message_2 = context.update(cx, |context, cx| {
3110 context
3111 .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)
3112 .unwrap()
3113 });
3114 buffer.update(cx, |buffer, cx| {
3115 buffer.edit([(0..0, "a"), (1..1, "b\nc")], None, cx);
3116 buffer.finalize_last_transaction();
3117 });
3118 let _message_3 = context.update(cx, |context, cx| {
3119 context
3120 .insert_message_after(message_2.id, Role::System, MessageStatus::Done, cx)
3121 .unwrap()
3122 });
3123 buffer.update(cx, |buffer, cx| buffer.undo(cx));
3124 assert_eq!(buffer.read_with(cx, |buffer, _| buffer.text()), "a\nb\nc\n");
3125 assert_eq!(
3126 cx.read(|cx| messages(&context, cx)),
3127 [
3128 (message_0, Role::User, 0..2),
3129 (message_1.id, Role::Assistant, 2..6),
3130 (message_2.id, Role::System, 6..6),
3131 ]
3132 );
3133
3134 let serialized_context = context.read_with(cx, |context, cx| context.serialize(cx));
3135 let deserialized_context = cx.new_model(|cx| {
3136 Context::deserialize(
3137 serialized_context,
3138 Default::default(),
3139 registry.clone(),
3140 prompt_builder.clone(),
3141 None,
3142 None,
3143 cx,
3144 )
3145 });
3146 let deserialized_buffer =
3147 deserialized_context.read_with(cx, |context, _| context.buffer.clone());
3148 assert_eq!(
3149 deserialized_buffer.read_with(cx, |buffer, _| buffer.text()),
3150 "a\nb\nc\n"
3151 );
3152 assert_eq!(
3153 cx.read(|cx| messages(&deserialized_context, cx)),
3154 [
3155 (message_0, Role::User, 0..2),
3156 (message_1.id, Role::Assistant, 2..6),
3157 (message_2.id, Role::System, 6..6),
3158 ]
3159 );
3160 }
3161
3162 #[gpui::test(iterations = 100)]
3163 async fn test_random_context_collaboration(cx: &mut TestAppContext, mut rng: StdRng) {
3164 let min_peers = env::var("MIN_PEERS")
3165 .map(|i| i.parse().expect("invalid `MIN_PEERS` variable"))
3166 .unwrap_or(2);
3167 let max_peers = env::var("MAX_PEERS")
3168 .map(|i| i.parse().expect("invalid `MAX_PEERS` variable"))
3169 .unwrap_or(5);
3170 let operations = env::var("OPERATIONS")
3171 .map(|i| i.parse().expect("invalid `OPERATIONS` variable"))
3172 .unwrap_or(50);
3173
3174 let settings_store = cx.update(SettingsStore::test);
3175 cx.set_global(settings_store);
3176 cx.update(LanguageModelRegistry::test);
3177
3178 cx.update(assistant_panel::init);
3179 let slash_commands = cx.update(SlashCommandRegistry::default_global);
3180 slash_commands.register_command(FakeSlashCommand("cmd-1".into()), false);
3181 slash_commands.register_command(FakeSlashCommand("cmd-2".into()), false);
3182 slash_commands.register_command(FakeSlashCommand("cmd-3".into()), false);
3183
3184 let registry = Arc::new(LanguageRegistry::test(cx.background_executor.clone()));
3185 let network = Arc::new(Mutex::new(Network::new(rng.clone())));
3186 let mut contexts = Vec::new();
3187
3188 let num_peers = rng.gen_range(min_peers..=max_peers);
3189 let context_id = ContextId::new();
3190 let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3191 for i in 0..num_peers {
3192 let context = cx.new_model(|cx| {
3193 Context::new(
3194 context_id.clone(),
3195 i as ReplicaId,
3196 language::Capability::ReadWrite,
3197 registry.clone(),
3198 prompt_builder.clone(),
3199 None,
3200 None,
3201 cx,
3202 )
3203 });
3204
3205 cx.update(|cx| {
3206 cx.subscribe(&context, {
3207 let network = network.clone();
3208 move |_, event, _| {
3209 if let ContextEvent::Operation(op) = event {
3210 network
3211 .lock()
3212 .broadcast(i as ReplicaId, vec![op.to_proto()]);
3213 }
3214 }
3215 })
3216 .detach();
3217 });
3218
3219 contexts.push(context);
3220 network.lock().add_peer(i as ReplicaId);
3221 }
3222
3223 let mut mutation_count = operations;
3224
3225 while mutation_count > 0
3226 || !network.lock().is_idle()
3227 || network.lock().contains_disconnected_peers()
3228 {
3229 let context_index = rng.gen_range(0..contexts.len());
3230 let context = &contexts[context_index];
3231
3232 match rng.gen_range(0..100) {
3233 0..=29 if mutation_count > 0 => {
3234 log::info!("Context {}: edit buffer", context_index);
3235 context.update(cx, |context, cx| {
3236 context
3237 .buffer
3238 .update(cx, |buffer, cx| buffer.randomly_edit(&mut rng, 1, cx));
3239 });
3240 mutation_count -= 1;
3241 }
3242 30..=44 if mutation_count > 0 => {
3243 context.update(cx, |context, cx| {
3244 let range = context.buffer.read(cx).random_byte_range(0, &mut rng);
3245 log::info!("Context {}: split message at {:?}", context_index, range);
3246 context.split_message(range, cx);
3247 });
3248 mutation_count -= 1;
3249 }
3250 45..=59 if mutation_count > 0 => {
3251 context.update(cx, |context, cx| {
3252 if let Some(message) = context.messages(cx).choose(&mut rng) {
3253 let role = *[Role::User, Role::Assistant, Role::System]
3254 .choose(&mut rng)
3255 .unwrap();
3256 log::info!(
3257 "Context {}: insert message after {:?} with {:?}",
3258 context_index,
3259 message.id,
3260 role
3261 );
3262 context.insert_message_after(message.id, role, MessageStatus::Done, cx);
3263 }
3264 });
3265 mutation_count -= 1;
3266 }
3267 60..=74 if mutation_count > 0 => {
3268 context.update(cx, |context, cx| {
3269 let command_text = "/".to_string()
3270 + slash_commands
3271 .command_names()
3272 .choose(&mut rng)
3273 .unwrap()
3274 .clone()
3275 .as_ref();
3276
3277 let command_range = context.buffer.update(cx, |buffer, cx| {
3278 let offset = buffer.random_byte_range(0, &mut rng).start;
3279 buffer.edit(
3280 [(offset..offset, format!("\n{}\n", command_text))],
3281 None,
3282 cx,
3283 );
3284 offset + 1..offset + 1 + command_text.len()
3285 });
3286
3287 let output_len = rng.gen_range(1..=10);
3288 let output_text = RandomCharIter::new(&mut rng)
3289 .filter(|c| *c != '\r')
3290 .take(output_len)
3291 .collect::<String>();
3292
3293 let num_sections = rng.gen_range(0..=3);
3294 let mut sections = Vec::with_capacity(num_sections);
3295 for _ in 0..num_sections {
3296 let section_start = rng.gen_range(0..output_len);
3297 let section_end = rng.gen_range(section_start..=output_len);
3298 sections.push(SlashCommandOutputSection {
3299 range: section_start..section_end,
3300 icon: ui::IconName::Ai,
3301 label: "section".into(),
3302 });
3303 }
3304
3305 log::info!(
3306 "Context {}: insert slash command output at {:?} with {:?}",
3307 context_index,
3308 command_range,
3309 sections
3310 );
3311
3312 let command_range =
3313 context.buffer.read(cx).anchor_after(command_range.start)
3314 ..context.buffer.read(cx).anchor_after(command_range.end);
3315 context.insert_command_output(
3316 command_range,
3317 Task::ready(Ok(SlashCommandOutput {
3318 text: output_text,
3319 sections,
3320 run_commands_in_text: false,
3321 })),
3322 true,
3323 cx,
3324 );
3325 });
3326 cx.run_until_parked();
3327 mutation_count -= 1;
3328 }
3329 75..=84 if mutation_count > 0 => {
3330 context.update(cx, |context, cx| {
3331 if let Some(message) = context.messages(cx).choose(&mut rng) {
3332 let new_status = match rng.gen_range(0..3) {
3333 0 => MessageStatus::Done,
3334 1 => MessageStatus::Pending,
3335 _ => MessageStatus::Error(SharedString::from("Random error")),
3336 };
3337 log::info!(
3338 "Context {}: update message {:?} status to {:?}",
3339 context_index,
3340 message.id,
3341 new_status
3342 );
3343 context.update_metadata(message.id, cx, |metadata| {
3344 metadata.status = new_status;
3345 });
3346 }
3347 });
3348 mutation_count -= 1;
3349 }
3350 _ => {
3351 let replica_id = context_index as ReplicaId;
3352 if network.lock().is_disconnected(replica_id) {
3353 network.lock().reconnect_peer(replica_id, 0);
3354
3355 let (ops_to_send, ops_to_receive) = cx.read(|cx| {
3356 let host_context = &contexts[0].read(cx);
3357 let guest_context = context.read(cx);
3358 (
3359 guest_context.serialize_ops(&host_context.version(cx), cx),
3360 host_context.serialize_ops(&guest_context.version(cx), cx),
3361 )
3362 });
3363 let ops_to_send = ops_to_send.await;
3364 let ops_to_receive = ops_to_receive
3365 .await
3366 .into_iter()
3367 .map(ContextOperation::from_proto)
3368 .collect::<Result<Vec<_>>>()
3369 .unwrap();
3370 log::info!(
3371 "Context {}: reconnecting. Sent {} operations, received {} operations",
3372 context_index,
3373 ops_to_send.len(),
3374 ops_to_receive.len()
3375 );
3376
3377 network.lock().broadcast(replica_id, ops_to_send);
3378 context
3379 .update(cx, |context, cx| context.apply_ops(ops_to_receive, cx))
3380 .unwrap();
3381 } else if rng.gen_bool(0.1) && replica_id != 0 {
3382 log::info!("Context {}: disconnecting", context_index);
3383 network.lock().disconnect_peer(replica_id);
3384 } else if network.lock().has_unreceived(replica_id) {
3385 log::info!("Context {}: applying operations", context_index);
3386 let ops = network.lock().receive(replica_id);
3387 let ops = ops
3388 .into_iter()
3389 .map(ContextOperation::from_proto)
3390 .collect::<Result<Vec<_>>>()
3391 .unwrap();
3392 context
3393 .update(cx, |context, cx| context.apply_ops(ops, cx))
3394 .unwrap();
3395 }
3396 }
3397 }
3398 }
3399
3400 cx.read(|cx| {
3401 let first_context = contexts[0].read(cx);
3402 for context in &contexts[1..] {
3403 let context = context.read(cx);
3404 assert!(context.pending_ops.is_empty());
3405 assert_eq!(
3406 context.buffer.read(cx).text(),
3407 first_context.buffer.read(cx).text(),
3408 "Context {} text != Context 0 text",
3409 context.buffer.read(cx).replica_id()
3410 );
3411 assert_eq!(
3412 context.message_anchors,
3413 first_context.message_anchors,
3414 "Context {} messages != Context 0 messages",
3415 context.buffer.read(cx).replica_id()
3416 );
3417 assert_eq!(
3418 context.messages_metadata,
3419 first_context.messages_metadata,
3420 "Context {} message metadata != Context 0 message metadata",
3421 context.buffer.read(cx).replica_id()
3422 );
3423 assert_eq!(
3424 context.slash_command_output_sections,
3425 first_context.slash_command_output_sections,
3426 "Context {} slash command output sections != Context 0 slash command output sections",
3427 context.buffer.read(cx).replica_id()
3428 );
3429 }
3430 });
3431 }
3432
3433 fn messages(context: &Model<Context>, cx: &AppContext) -> Vec<(MessageId, Role, Range<usize>)> {
3434 context
3435 .read(cx)
3436 .messages(cx)
3437 .map(|message| (message.id, message.role, message.offset_range))
3438 .collect()
3439 }
3440
3441 #[derive(Clone)]
3442 struct FakeSlashCommand(String);
3443
3444 impl SlashCommand for FakeSlashCommand {
3445 fn name(&self) -> String {
3446 self.0.clone()
3447 }
3448
3449 fn description(&self) -> String {
3450 format!("Fake slash command: {}", self.0)
3451 }
3452
3453 fn menu_text(&self) -> String {
3454 format!("Run fake command: {}", self.0)
3455 }
3456
3457 fn complete_argument(
3458 self: Arc<Self>,
3459 _query: String,
3460 _cancel: Arc<AtomicBool>,
3461 _workspace: Option<WeakView<Workspace>>,
3462 _cx: &mut AppContext,
3463 ) -> Task<Result<Vec<ArgumentCompletion>>> {
3464 Task::ready(Ok(vec![]))
3465 }
3466
3467 fn requires_argument(&self) -> bool {
3468 false
3469 }
3470
3471 fn run(
3472 self: Arc<Self>,
3473 _argument: Option<&str>,
3474 _workspace: WeakView<Workspace>,
3475 _delegate: Option<Arc<dyn LspAdapterDelegate>>,
3476 _cx: &mut WindowContext,
3477 ) -> Task<Result<SlashCommandOutput>> {
3478 Task::ready(Ok(SlashCommandOutput {
3479 text: format!("Executed fake command: {}", self.0),
3480 sections: vec![],
3481 run_commands_in_text: false,
3482 }))
3483 }
3484 }
3485}
3486
3487mod tool {
3488 use gpui::AsyncAppContext;
3489
3490 use super::*;
3491
3492 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
3493 pub struct WorkflowStepEditSuggestions {
3494 /// An extremely short title for the edit step represented by these operations.
3495 pub step_title: String,
3496 /// A sequence of operations to apply to the codebase.
3497 /// When multiple operations are required for a step, be sure to include multiple operations in this list.
3498 pub edit_suggestions: Vec<EditSuggestion>,
3499 }
3500
3501 impl LanguageModelTool for WorkflowStepEditSuggestions {
3502 fn name() -> String {
3503 "edit".into()
3504 }
3505
3506 fn description() -> String {
3507 "suggest edits to one or more locations in the codebase".into()
3508 }
3509 }
3510
3511 /// A description of an operation to apply to one location in the codebase.
3512 ///
3513 /// This object represents a single edit operation that can be performed on a specific file
3514 /// in the codebase. It encapsulates both the location (file path) and the nature of the
3515 /// edit to be made.
3516 ///
3517 /// # Fields
3518 ///
3519 /// * `path`: A string representing the file path where the edit operation should be applied.
3520 /// This path is relative to the root of the project or repository.
3521 ///
3522 /// * `kind`: An enum representing the specific type of edit operation to be performed.
3523 ///
3524 /// # Usage
3525 ///
3526 /// `EditOperation` is used within a code editor to represent and apply
3527 /// programmatic changes to source code. It provides a structured way to describe
3528 /// edits for features like refactoring tools or AI-assisted coding suggestions.
3529 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
3530 pub struct EditSuggestion {
3531 /// The path to the file containing the relevant operation
3532 pub path: String,
3533 #[serde(flatten)]
3534 pub kind: EditSuggestionKind,
3535 }
3536
3537 impl EditSuggestion {
3538 pub(super) async fn resolve(
3539 &self,
3540 project: Model<Project>,
3541 mut cx: AsyncAppContext,
3542 ) -> Result<(Model<Buffer>, super::EditSuggestion)> {
3543 let path = self.path.clone();
3544 let kind = self.kind.clone();
3545 let buffer = project
3546 .update(&mut cx, |project, cx| {
3547 let project_path = project
3548 .find_project_path(Path::new(&path), cx)
3549 .with_context(|| format!("worktree not found for {:?}", path))?;
3550 anyhow::Ok(project.open_buffer(project_path, cx))
3551 })??
3552 .await?;
3553
3554 let mut parse_status = buffer.read_with(&cx, |buffer, _cx| buffer.parse_status())?;
3555 while *parse_status.borrow() != ParseStatus::Idle {
3556 parse_status.changed().await?;
3557 }
3558
3559 let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
3560 let outline = snapshot.outline(None).context("no outline for buffer")?;
3561
3562 let suggestion;
3563 match kind {
3564 EditSuggestionKind::Update {
3565 symbol,
3566 description,
3567 } => {
3568 let symbol = outline
3569 .find_most_similar(&symbol)
3570 .with_context(|| format!("symbol not found: {:?}", symbol))?
3571 .to_point(&snapshot);
3572 let start = symbol
3573 .annotation_range
3574 .map_or(symbol.range.start, |range| range.start);
3575 let start = Point::new(start.row, 0);
3576 let end = Point::new(
3577 symbol.range.end.row,
3578 snapshot.line_len(symbol.range.end.row),
3579 );
3580 let range = snapshot.anchor_before(start)..snapshot.anchor_after(end);
3581 suggestion = super::EditSuggestion::Update { range, description };
3582 }
3583 EditSuggestionKind::Create { description } => {
3584 suggestion = super::EditSuggestion::CreateFile { description };
3585 }
3586 EditSuggestionKind::InsertSiblingBefore {
3587 symbol,
3588 description,
3589 } => {
3590 let symbol = outline
3591 .find_most_similar(&symbol)
3592 .with_context(|| format!("symbol not found: {:?}", symbol))?
3593 .to_point(&snapshot);
3594 let position = snapshot.anchor_before(
3595 symbol
3596 .annotation_range
3597 .map_or(symbol.range.start, |annotation_range| {
3598 annotation_range.start
3599 }),
3600 );
3601 suggestion = super::EditSuggestion::InsertSiblingBefore {
3602 position,
3603 description,
3604 };
3605 }
3606 EditSuggestionKind::InsertSiblingAfter {
3607 symbol,
3608 description,
3609 } => {
3610 let symbol = outline
3611 .find_most_similar(&symbol)
3612 .with_context(|| format!("symbol not found: {:?}", symbol))?
3613 .to_point(&snapshot);
3614 let position = snapshot.anchor_after(symbol.range.end);
3615 suggestion = super::EditSuggestion::InsertSiblingAfter {
3616 position,
3617 description,
3618 };
3619 }
3620 EditSuggestionKind::PrependChild {
3621 symbol,
3622 description,
3623 } => {
3624 if let Some(symbol) = symbol {
3625 let symbol = outline
3626 .find_most_similar(&symbol)
3627 .with_context(|| format!("symbol not found: {:?}", symbol))?
3628 .to_point(&snapshot);
3629
3630 let position = snapshot.anchor_after(
3631 symbol
3632 .body_range
3633 .map_or(symbol.range.start, |body_range| body_range.start),
3634 );
3635 suggestion = super::EditSuggestion::PrependChild {
3636 position,
3637 description,
3638 };
3639 } else {
3640 suggestion = super::EditSuggestion::PrependChild {
3641 position: language::Anchor::MIN,
3642 description,
3643 };
3644 }
3645 }
3646 EditSuggestionKind::AppendChild {
3647 symbol,
3648 description,
3649 } => {
3650 if let Some(symbol) = symbol {
3651 let symbol = outline
3652 .find_most_similar(&symbol)
3653 .with_context(|| format!("symbol not found: {:?}", symbol))?
3654 .to_point(&snapshot);
3655
3656 let position = snapshot.anchor_before(
3657 symbol
3658 .body_range
3659 .map_or(symbol.range.end, |body_range| body_range.end),
3660 );
3661 suggestion = super::EditSuggestion::AppendChild {
3662 position,
3663 description,
3664 };
3665 } else {
3666 suggestion = super::EditSuggestion::PrependChild {
3667 position: language::Anchor::MAX,
3668 description,
3669 };
3670 }
3671 }
3672 EditSuggestionKind::Delete { symbol } => {
3673 let symbol = outline
3674 .find_most_similar(&symbol)
3675 .with_context(|| format!("symbol not found: {:?}", symbol))?
3676 .to_point(&snapshot);
3677 let start = symbol
3678 .annotation_range
3679 .map_or(symbol.range.start, |range| range.start);
3680 let start = Point::new(start.row, 0);
3681 let end = Point::new(
3682 symbol.range.end.row,
3683 snapshot.line_len(symbol.range.end.row),
3684 );
3685 let range = snapshot.anchor_before(start)..snapshot.anchor_after(end);
3686 suggestion = super::EditSuggestion::Delete { range };
3687 }
3688 }
3689
3690 Ok((buffer, suggestion))
3691 }
3692 }
3693
3694 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
3695 #[serde(tag = "kind")]
3696 pub enum EditSuggestionKind {
3697 /// Rewrites the specified symbol entirely based on the given description.
3698 /// This operation completely replaces the existing symbol with new content.
3699 Update {
3700 /// A fully-qualified reference to the symbol, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
3701 /// The path should uniquely identify the symbol within the containing file.
3702 symbol: String,
3703 /// A brief description of the transformation to apply to the symbol.
3704 description: String,
3705 },
3706 /// Creates a new file with the given path based on the provided description.
3707 /// This operation adds a new file to the codebase.
3708 Create {
3709 /// A brief description of the file to be created.
3710 description: String,
3711 },
3712 /// Inserts a new symbol based on the given description before the specified symbol.
3713 /// This operation adds new content immediately preceding an existing symbol.
3714 InsertSiblingBefore {
3715 /// A fully-qualified reference to the symbol, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
3716 /// The new content will be inserted immediately before this symbol.
3717 symbol: String,
3718 /// A brief description of the new symbol to be inserted.
3719 description: String,
3720 },
3721 /// Inserts a new symbol based on the given description after the specified symbol.
3722 /// This operation adds new content immediately following an existing symbol.
3723 InsertSiblingAfter {
3724 /// A fully-qualified reference to the symbol, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
3725 /// The new content will be inserted immediately after this symbol.
3726 symbol: String,
3727 /// A brief description of the new symbol to be inserted.
3728 description: String,
3729 },
3730 /// Inserts a new symbol as a child of the specified symbol at the start.
3731 /// This operation adds new content as the first child of an existing symbol (or file if no symbol is provided).
3732 PrependChild {
3733 /// 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`.
3734 /// If provided, the new content will be inserted as the first child of this symbol.
3735 /// If not provided, the new content will be inserted at the top of the file.
3736 symbol: Option<String>,
3737 /// A brief description of the new symbol to be inserted.
3738 description: String,
3739 },
3740 /// Inserts a new symbol as a child of the specified symbol at the end.
3741 /// This operation adds new content as the last child of an existing symbol (or file if no symbol is provided).
3742 AppendChild {
3743 /// 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`.
3744 /// If provided, the new content will be inserted as the last child of this symbol.
3745 /// If not provided, the new content will be applied at the bottom of the file.
3746 symbol: Option<String>,
3747 /// A brief description of the new symbol to be inserted.
3748 description: String,
3749 },
3750 /// Deletes the specified symbol from the containing file.
3751 Delete {
3752 /// An fully-qualified reference to the symbol to be deleted, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
3753 symbol: String,
3754 },
3755 }
3756
3757 impl EditSuggestionKind {
3758 pub fn symbol(&self) -> Option<&str> {
3759 match self {
3760 Self::Update { symbol, .. } => Some(symbol),
3761 Self::InsertSiblingBefore { symbol, .. } => Some(symbol),
3762 Self::InsertSiblingAfter { symbol, .. } => Some(symbol),
3763 Self::PrependChild { symbol, .. } => symbol.as_deref(),
3764 Self::AppendChild { symbol, .. } => symbol.as_deref(),
3765 Self::Delete { symbol } => Some(symbol),
3766 Self::Create { .. } => None,
3767 }
3768 }
3769
3770 pub fn description(&self) -> Option<&str> {
3771 match self {
3772 Self::Update { description, .. } => Some(description),
3773 Self::Create { description } => Some(description),
3774 Self::InsertSiblingBefore { description, .. } => Some(description),
3775 Self::InsertSiblingAfter { description, .. } => Some(description),
3776 Self::PrependChild { description, .. } => Some(description),
3777 Self::AppendChild { description, .. } => Some(description),
3778 Self::Delete { .. } => None,
3779 }
3780 }
3781
3782 pub fn initial_insertion(&self) -> Option<InitialInsertion> {
3783 match self {
3784 EditSuggestionKind::InsertSiblingBefore { .. } => {
3785 Some(InitialInsertion::NewlineAfter)
3786 }
3787 EditSuggestionKind::InsertSiblingAfter { .. } => {
3788 Some(InitialInsertion::NewlineBefore)
3789 }
3790 EditSuggestionKind::PrependChild { .. } => Some(InitialInsertion::NewlineAfter),
3791 EditSuggestionKind::AppendChild { .. } => Some(InitialInsertion::NewlineBefore),
3792 _ => None,
3793 }
3794 }
3795 }
3796}