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