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