context.rs

   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| &section.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}