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