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