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