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