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