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