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