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