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