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