context.rs

   1#[cfg(test)]
   2mod context_tests;
   3
   4use agent_settings::AgentSettings;
   5use anyhow::{Context as _, Result, bail};
   6use assistant_slash_command::{
   7    SlashCommandContent, SlashCommandEvent, SlashCommandLine, SlashCommandOutputSection,
   8    SlashCommandResult, SlashCommandWorkingSet,
   9};
  10use assistant_slash_commands::FileCommandMetadata;
  11use client::{self, proto, telemetry::Telemetry};
  12use clock::ReplicaId;
  13use collections::{HashMap, HashSet};
  14use fs::{Fs, RemoveOptions};
  15use futures::{FutureExt, StreamExt, future::Shared};
  16use gpui::{
  17    App, AppContext as _, Context, Entity, EventEmitter, RenderImage, SharedString, Subscription,
  18    Task,
  19};
  20use language::{AnchorRangeExt, Bias, Buffer, LanguageRegistry, OffsetRangeExt, Point, ToOffset};
  21use language_model::{
  22    LanguageModel, LanguageModelCacheConfiguration, LanguageModelCompletionEvent,
  23    LanguageModelImage, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
  24    LanguageModelToolUseId, MessageContent, PaymentRequiredError, Role, StopReason,
  25    report_assistant_event,
  26};
  27use open_ai::Model as OpenAiModel;
  28use paths::contexts_dir;
  29use project::Project;
  30use prompt_store::PromptBuilder;
  31use serde::{Deserialize, Serialize};
  32use settings::Settings;
  33use smallvec::SmallVec;
  34use std::{
  35    cmp::{Ordering, max},
  36    fmt::{Debug, Write as _},
  37    iter, mem,
  38    ops::Range,
  39    path::Path,
  40    sync::Arc,
  41    time::{Duration, Instant},
  42};
  43use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};
  44use text::{BufferSnapshot, ToPoint};
  45use ui::IconName;
  46use util::{ResultExt, TryFutureExt, post_inc};
  47use uuid::Uuid;
  48
  49#[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
  50pub struct ContextId(String);
  51
  52impl ContextId {
  53    pub fn new() -> Self {
  54        Self(Uuid::new_v4().to_string())
  55    }
  56
  57    pub fn from_proto(id: String) -> Self {
  58        Self(id)
  59    }
  60
  61    pub fn to_proto(&self) -> String {
  62        self.0.clone()
  63    }
  64}
  65
  66#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
  67pub struct MessageId(pub clock::Lamport);
  68
  69impl MessageId {
  70    pub fn as_u64(self) -> u64 {
  71        self.0.as_u64()
  72    }
  73}
  74
  75#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
  76pub enum MessageStatus {
  77    Pending,
  78    Done,
  79    Error(SharedString),
  80    Canceled,
  81}
  82
  83impl MessageStatus {
  84    pub fn from_proto(status: proto::ContextMessageStatus) -> MessageStatus {
  85        match status.variant {
  86            Some(proto::context_message_status::Variant::Pending(_)) => MessageStatus::Pending,
  87            Some(proto::context_message_status::Variant::Done(_)) => MessageStatus::Done,
  88            Some(proto::context_message_status::Variant::Error(error)) => {
  89                MessageStatus::Error(error.message.into())
  90            }
  91            Some(proto::context_message_status::Variant::Canceled(_)) => MessageStatus::Canceled,
  92            None => MessageStatus::Pending,
  93        }
  94    }
  95
  96    pub fn to_proto(&self) -> proto::ContextMessageStatus {
  97        match self {
  98            MessageStatus::Pending => proto::ContextMessageStatus {
  99                variant: Some(proto::context_message_status::Variant::Pending(
 100                    proto::context_message_status::Pending {},
 101                )),
 102            },
 103            MessageStatus::Done => proto::ContextMessageStatus {
 104                variant: Some(proto::context_message_status::Variant::Done(
 105                    proto::context_message_status::Done {},
 106                )),
 107            },
 108            MessageStatus::Error(message) => proto::ContextMessageStatus {
 109                variant: Some(proto::context_message_status::Variant::Error(
 110                    proto::context_message_status::Error {
 111                        message: message.to_string(),
 112                    },
 113                )),
 114            },
 115            MessageStatus::Canceled => proto::ContextMessageStatus {
 116                variant: Some(proto::context_message_status::Variant::Canceled(
 117                    proto::context_message_status::Canceled {},
 118                )),
 119            },
 120        }
 121    }
 122}
 123
 124#[derive(Clone, Debug)]
 125pub enum ContextOperation {
 126    InsertMessage {
 127        anchor: MessageAnchor,
 128        metadata: MessageMetadata,
 129        version: clock::Global,
 130    },
 131    UpdateMessage {
 132        message_id: MessageId,
 133        metadata: MessageMetadata,
 134        version: clock::Global,
 135    },
 136    UpdateSummary {
 137        summary: ContextSummaryContent,
 138        version: clock::Global,
 139    },
 140    SlashCommandStarted {
 141        id: InvokedSlashCommandId,
 142        output_range: Range<language::Anchor>,
 143        name: String,
 144        version: clock::Global,
 145    },
 146    SlashCommandFinished {
 147        id: InvokedSlashCommandId,
 148        timestamp: clock::Lamport,
 149        error_message: Option<String>,
 150        version: clock::Global,
 151    },
 152    SlashCommandOutputSectionAdded {
 153        timestamp: clock::Lamport,
 154        section: SlashCommandOutputSection<language::Anchor>,
 155        version: clock::Global,
 156    },
 157    ThoughtProcessOutputSectionAdded {
 158        timestamp: clock::Lamport,
 159        section: ThoughtProcessOutputSection<language::Anchor>,
 160        version: clock::Global,
 161    },
 162    BufferOperation(language::Operation),
 163}
 164
 165impl ContextOperation {
 166    pub fn from_proto(op: proto::ContextOperation) -> Result<Self> {
 167        match op.variant.context("invalid variant")? {
 168            proto::context_operation::Variant::InsertMessage(insert) => {
 169                let message = insert.message.context("invalid message")?;
 170                let id = MessageId(language::proto::deserialize_timestamp(
 171                    message.id.context("invalid id")?,
 172                ));
 173                Ok(Self::InsertMessage {
 174                    anchor: MessageAnchor {
 175                        id,
 176                        start: language::proto::deserialize_anchor(
 177                            message.start.context("invalid anchor")?,
 178                        )
 179                        .context("invalid anchor")?,
 180                    },
 181                    metadata: MessageMetadata {
 182                        role: Role::from_proto(message.role),
 183                        status: MessageStatus::from_proto(
 184                            message.status.context("invalid status")?,
 185                        ),
 186                        timestamp: id.0,
 187                        cache: None,
 188                    },
 189                    version: language::proto::deserialize_version(&insert.version),
 190                })
 191            }
 192            proto::context_operation::Variant::UpdateMessage(update) => Ok(Self::UpdateMessage {
 193                message_id: MessageId(language::proto::deserialize_timestamp(
 194                    update.message_id.context("invalid message id")?,
 195                )),
 196                metadata: MessageMetadata {
 197                    role: Role::from_proto(update.role),
 198                    status: MessageStatus::from_proto(update.status.context("invalid status")?),
 199                    timestamp: language::proto::deserialize_timestamp(
 200                        update.timestamp.context("invalid timestamp")?,
 201                    ),
 202                    cache: None,
 203                },
 204                version: language::proto::deserialize_version(&update.version),
 205            }),
 206            proto::context_operation::Variant::UpdateSummary(update) => Ok(Self::UpdateSummary {
 207                summary: ContextSummaryContent {
 208                    text: update.summary,
 209                    done: update.done,
 210                    timestamp: language::proto::deserialize_timestamp(
 211                        update.timestamp.context("invalid timestamp")?,
 212                    ),
 213                },
 214                version: language::proto::deserialize_version(&update.version),
 215            }),
 216            proto::context_operation::Variant::SlashCommandStarted(message) => {
 217                Ok(Self::SlashCommandStarted {
 218                    id: InvokedSlashCommandId(language::proto::deserialize_timestamp(
 219                        message.id.context("invalid id")?,
 220                    )),
 221                    output_range: language::proto::deserialize_anchor_range(
 222                        message.output_range.context("invalid range")?,
 223                    )?,
 224                    name: message.name,
 225                    version: language::proto::deserialize_version(&message.version),
 226                })
 227            }
 228            proto::context_operation::Variant::SlashCommandOutputSectionAdded(message) => {
 229                let section = message.section.context("missing section")?;
 230                Ok(Self::SlashCommandOutputSectionAdded {
 231                    timestamp: language::proto::deserialize_timestamp(
 232                        message.timestamp.context("missing timestamp")?,
 233                    ),
 234                    section: SlashCommandOutputSection {
 235                        range: language::proto::deserialize_anchor_range(
 236                            section.range.context("invalid range")?,
 237                        )?,
 238                        icon: section.icon_name.parse()?,
 239                        label: section.label.into(),
 240                        metadata: section
 241                            .metadata
 242                            .and_then(|metadata| serde_json::from_str(&metadata).log_err()),
 243                    },
 244                    version: language::proto::deserialize_version(&message.version),
 245                })
 246            }
 247            proto::context_operation::Variant::SlashCommandCompleted(message) => {
 248                Ok(Self::SlashCommandFinished {
 249                    id: InvokedSlashCommandId(language::proto::deserialize_timestamp(
 250                        message.id.context("invalid id")?,
 251                    )),
 252                    timestamp: language::proto::deserialize_timestamp(
 253                        message.timestamp.context("missing timestamp")?,
 254                    ),
 255                    error_message: message.error_message,
 256                    version: language::proto::deserialize_version(&message.version),
 257                })
 258            }
 259            proto::context_operation::Variant::ThoughtProcessOutputSectionAdded(message) => {
 260                let section = message.section.context("missing section")?;
 261                Ok(Self::ThoughtProcessOutputSectionAdded {
 262                    timestamp: language::proto::deserialize_timestamp(
 263                        message.timestamp.context("missing timestamp")?,
 264                    ),
 265                    section: ThoughtProcessOutputSection {
 266                        range: language::proto::deserialize_anchor_range(
 267                            section.range.context("invalid range")?,
 268                        )?,
 269                    },
 270                    version: language::proto::deserialize_version(&message.version),
 271                })
 272            }
 273            proto::context_operation::Variant::BufferOperation(op) => Ok(Self::BufferOperation(
 274                language::proto::deserialize_operation(
 275                    op.operation.context("invalid buffer operation")?,
 276                )?,
 277            )),
 278        }
 279    }
 280
 281    pub fn to_proto(&self) -> proto::ContextOperation {
 282        match self {
 283            Self::InsertMessage {
 284                anchor,
 285                metadata,
 286                version,
 287            } => proto::ContextOperation {
 288                variant: Some(proto::context_operation::Variant::InsertMessage(
 289                    proto::context_operation::InsertMessage {
 290                        message: Some(proto::ContextMessage {
 291                            id: Some(language::proto::serialize_timestamp(anchor.id.0)),
 292                            start: Some(language::proto::serialize_anchor(&anchor.start)),
 293                            role: metadata.role.to_proto() as i32,
 294                            status: Some(metadata.status.to_proto()),
 295                        }),
 296                        version: language::proto::serialize_version(version),
 297                    },
 298                )),
 299            },
 300            Self::UpdateMessage {
 301                message_id,
 302                metadata,
 303                version,
 304            } => proto::ContextOperation {
 305                variant: Some(proto::context_operation::Variant::UpdateMessage(
 306                    proto::context_operation::UpdateMessage {
 307                        message_id: Some(language::proto::serialize_timestamp(message_id.0)),
 308                        role: metadata.role.to_proto() as i32,
 309                        status: Some(metadata.status.to_proto()),
 310                        timestamp: Some(language::proto::serialize_timestamp(metadata.timestamp)),
 311                        version: language::proto::serialize_version(version),
 312                    },
 313                )),
 314            },
 315            Self::UpdateSummary { summary, version } => proto::ContextOperation {
 316                variant: Some(proto::context_operation::Variant::UpdateSummary(
 317                    proto::context_operation::UpdateSummary {
 318                        summary: summary.text.clone(),
 319                        done: summary.done,
 320                        timestamp: Some(language::proto::serialize_timestamp(summary.timestamp)),
 321                        version: language::proto::serialize_version(version),
 322                    },
 323                )),
 324            },
 325            Self::SlashCommandStarted {
 326                id,
 327                output_range,
 328                name,
 329                version,
 330            } => proto::ContextOperation {
 331                variant: Some(proto::context_operation::Variant::SlashCommandStarted(
 332                    proto::context_operation::SlashCommandStarted {
 333                        id: Some(language::proto::serialize_timestamp(id.0)),
 334                        output_range: Some(language::proto::serialize_anchor_range(
 335                            output_range.clone(),
 336                        )),
 337                        name: name.clone(),
 338                        version: language::proto::serialize_version(version),
 339                    },
 340                )),
 341            },
 342            Self::SlashCommandOutputSectionAdded {
 343                timestamp,
 344                section,
 345                version,
 346            } => proto::ContextOperation {
 347                variant: Some(
 348                    proto::context_operation::Variant::SlashCommandOutputSectionAdded(
 349                        proto::context_operation::SlashCommandOutputSectionAdded {
 350                            timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
 351                            section: Some({
 352                                let icon_name: &'static str = section.icon.into();
 353                                proto::SlashCommandOutputSection {
 354                                    range: Some(language::proto::serialize_anchor_range(
 355                                        section.range.clone(),
 356                                    )),
 357                                    icon_name: icon_name.to_string(),
 358                                    label: section.label.to_string(),
 359                                    metadata: section.metadata.as_ref().and_then(|metadata| {
 360                                        serde_json::to_string(metadata).log_err()
 361                                    }),
 362                                }
 363                            }),
 364                            version: language::proto::serialize_version(version),
 365                        },
 366                    ),
 367                ),
 368            },
 369            Self::SlashCommandFinished {
 370                id,
 371                timestamp,
 372                error_message,
 373                version,
 374            } => proto::ContextOperation {
 375                variant: Some(proto::context_operation::Variant::SlashCommandCompleted(
 376                    proto::context_operation::SlashCommandCompleted {
 377                        id: Some(language::proto::serialize_timestamp(id.0)),
 378                        timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
 379                        error_message: error_message.clone(),
 380                        version: language::proto::serialize_version(version),
 381                    },
 382                )),
 383            },
 384            Self::ThoughtProcessOutputSectionAdded {
 385                timestamp,
 386                section,
 387                version,
 388            } => proto::ContextOperation {
 389                variant: Some(
 390                    proto::context_operation::Variant::ThoughtProcessOutputSectionAdded(
 391                        proto::context_operation::ThoughtProcessOutputSectionAdded {
 392                            timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
 393                            section: Some({
 394                                proto::ThoughtProcessOutputSection {
 395                                    range: Some(language::proto::serialize_anchor_range(
 396                                        section.range.clone(),
 397                                    )),
 398                                }
 399                            }),
 400                            version: language::proto::serialize_version(version),
 401                        },
 402                    ),
 403                ),
 404            },
 405            Self::BufferOperation(operation) => proto::ContextOperation {
 406                variant: Some(proto::context_operation::Variant::BufferOperation(
 407                    proto::context_operation::BufferOperation {
 408                        operation: Some(language::proto::serialize_operation(operation)),
 409                    },
 410                )),
 411            },
 412        }
 413    }
 414
 415    fn timestamp(&self) -> clock::Lamport {
 416        match self {
 417            Self::InsertMessage { anchor, .. } => anchor.id.0,
 418            Self::UpdateMessage { metadata, .. } => metadata.timestamp,
 419            Self::UpdateSummary { summary, .. } => summary.timestamp,
 420            Self::SlashCommandStarted { id, .. } => id.0,
 421            Self::SlashCommandOutputSectionAdded { timestamp, .. }
 422            | Self::SlashCommandFinished { timestamp, .. }
 423            | Self::ThoughtProcessOutputSectionAdded { timestamp, .. } => *timestamp,
 424            Self::BufferOperation(_) => {
 425                panic!("reading the timestamp of a buffer operation is not supported")
 426            }
 427        }
 428    }
 429
 430    /// Returns the current version of the context operation.
 431    pub fn version(&self) -> &clock::Global {
 432        match self {
 433            Self::InsertMessage { version, .. }
 434            | Self::UpdateMessage { version, .. }
 435            | Self::UpdateSummary { version, .. }
 436            | Self::SlashCommandStarted { version, .. }
 437            | Self::SlashCommandOutputSectionAdded { version, .. }
 438            | Self::SlashCommandFinished { version, .. }
 439            | Self::ThoughtProcessOutputSectionAdded { version, .. } => version,
 440            Self::BufferOperation(_) => {
 441                panic!("reading the version of a buffer operation is not supported")
 442            }
 443        }
 444    }
 445}
 446
 447#[derive(Debug, Clone)]
 448pub enum ContextEvent {
 449    ShowAssistError(SharedString),
 450    ShowPaymentRequiredError,
 451    MessagesEdited,
 452    SummaryChanged,
 453    SummaryGenerated,
 454    StreamedCompletion,
 455    StartedThoughtProcess(Range<language::Anchor>),
 456    EndedThoughtProcess(language::Anchor),
 457    InvokedSlashCommandChanged {
 458        command_id: InvokedSlashCommandId,
 459    },
 460    ParsedSlashCommandsUpdated {
 461        removed: Vec<Range<language::Anchor>>,
 462        updated: Vec<ParsedSlashCommand>,
 463    },
 464    SlashCommandOutputSectionAdded {
 465        section: SlashCommandOutputSection<language::Anchor>,
 466    },
 467    Operation(ContextOperation),
 468}
 469
 470#[derive(Clone, Debug, Eq, PartialEq)]
 471pub enum ContextSummary {
 472    Pending,
 473    Content(ContextSummaryContent),
 474    Error,
 475}
 476
 477#[derive(Default, Clone, Debug, Eq, PartialEq)]
 478pub struct ContextSummaryContent {
 479    pub text: String,
 480    pub done: bool,
 481    pub timestamp: clock::Lamport,
 482}
 483
 484impl ContextSummary {
 485    pub const DEFAULT: &str = "New Text Thread";
 486
 487    pub fn or_default(&self) -> SharedString {
 488        self.unwrap_or(Self::DEFAULT)
 489    }
 490
 491    pub fn unwrap_or(&self, message: impl Into<SharedString>) -> SharedString {
 492        self.content()
 493            .map_or_else(|| message.into(), |content| content.text.clone().into())
 494    }
 495
 496    pub fn content(&self) -> Option<&ContextSummaryContent> {
 497        match self {
 498            ContextSummary::Content(content) => Some(content),
 499            ContextSummary::Pending | ContextSummary::Error => None,
 500        }
 501    }
 502
 503    fn content_as_mut(&mut self) -> Option<&mut ContextSummaryContent> {
 504        match self {
 505            ContextSummary::Content(content) => Some(content),
 506            ContextSummary::Pending | ContextSummary::Error => None,
 507        }
 508    }
 509
 510    fn content_or_set_empty(&mut self) -> &mut ContextSummaryContent {
 511        match self {
 512            ContextSummary::Content(content) => content,
 513            ContextSummary::Pending | ContextSummary::Error => {
 514                let content = ContextSummaryContent::default();
 515                *self = ContextSummary::Content(content);
 516                self.content_as_mut().unwrap()
 517            }
 518        }
 519    }
 520
 521    pub fn is_pending(&self) -> bool {
 522        matches!(self, ContextSummary::Pending)
 523    }
 524
 525    fn timestamp(&self) -> Option<clock::Lamport> {
 526        match self {
 527            ContextSummary::Content(content) => Some(content.timestamp),
 528            ContextSummary::Pending | ContextSummary::Error => None,
 529        }
 530    }
 531}
 532
 533impl PartialOrd for ContextSummary {
 534    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
 535        self.timestamp().partial_cmp(&other.timestamp())
 536    }
 537}
 538
 539#[derive(Clone, Debug, Eq, PartialEq)]
 540pub struct MessageAnchor {
 541    pub id: MessageId,
 542    pub start: language::Anchor,
 543}
 544
 545#[derive(Clone, Debug, Eq, PartialEq)]
 546pub enum CacheStatus {
 547    Pending,
 548    Cached,
 549}
 550
 551#[derive(Clone, Debug, Eq, PartialEq)]
 552pub struct MessageCacheMetadata {
 553    pub is_anchor: bool,
 554    pub is_final_anchor: bool,
 555    pub status: CacheStatus,
 556    pub cached_at: clock::Global,
 557}
 558
 559#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 560pub struct MessageMetadata {
 561    pub role: Role,
 562    pub status: MessageStatus,
 563    pub timestamp: clock::Lamport,
 564    #[serde(skip)]
 565    pub cache: Option<MessageCacheMetadata>,
 566}
 567
 568impl From<&Message> for MessageMetadata {
 569    fn from(message: &Message) -> Self {
 570        Self {
 571            role: message.role,
 572            status: message.status.clone(),
 573            timestamp: message.id.0,
 574            cache: message.cache.clone(),
 575        }
 576    }
 577}
 578
 579impl MessageMetadata {
 580    pub fn is_cache_valid(&self, buffer: &BufferSnapshot, range: &Range<usize>) -> bool {
 581        let result = match &self.cache {
 582            Some(MessageCacheMetadata { cached_at, .. }) => !buffer.has_edits_since_in_range(
 583                &cached_at,
 584                Range {
 585                    start: buffer.anchor_at(range.start, Bias::Right),
 586                    end: buffer.anchor_at(range.end, Bias::Left),
 587                },
 588            ),
 589            _ => false,
 590        };
 591        result
 592    }
 593}
 594
 595#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 596pub struct ThoughtProcessOutputSection<T> {
 597    pub range: Range<T>,
 598}
 599
 600impl ThoughtProcessOutputSection<language::Anchor> {
 601    pub fn is_valid(&self, buffer: &language::TextBuffer) -> bool {
 602        self.range.start.is_valid(buffer) && !self.range.to_offset(buffer).is_empty()
 603    }
 604}
 605
 606#[derive(Clone, Debug)]
 607pub struct Message {
 608    pub offset_range: Range<usize>,
 609    pub index_range: Range<usize>,
 610    pub anchor_range: Range<language::Anchor>,
 611    pub id: MessageId,
 612    pub role: Role,
 613    pub status: MessageStatus,
 614    pub cache: Option<MessageCacheMetadata>,
 615}
 616
 617#[derive(Debug, Clone)]
 618pub enum Content {
 619    Image {
 620        anchor: language::Anchor,
 621        image_id: u64,
 622        render_image: Arc<RenderImage>,
 623        image: Shared<Task<Option<LanguageModelImage>>>,
 624    },
 625}
 626
 627impl Content {
 628    fn range(&self) -> Range<language::Anchor> {
 629        match self {
 630            Self::Image { anchor, .. } => *anchor..*anchor,
 631        }
 632    }
 633
 634    fn cmp(&self, other: &Self, buffer: &BufferSnapshot) -> Ordering {
 635        let self_range = self.range();
 636        let other_range = other.range();
 637        if self_range.end.cmp(&other_range.start, buffer).is_lt() {
 638            Ordering::Less
 639        } else if self_range.start.cmp(&other_range.end, buffer).is_gt() {
 640            Ordering::Greater
 641        } else {
 642            Ordering::Equal
 643        }
 644    }
 645}
 646
 647struct PendingCompletion {
 648    id: usize,
 649    assistant_message_id: MessageId,
 650    _task: Task<()>,
 651}
 652
 653#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
 654pub struct InvokedSlashCommandId(clock::Lamport);
 655
 656pub struct AssistantContext {
 657    id: ContextId,
 658    timestamp: clock::Lamport,
 659    version: clock::Global,
 660    pending_ops: Vec<ContextOperation>,
 661    operations: Vec<ContextOperation>,
 662    buffer: Entity<Buffer>,
 663    parsed_slash_commands: Vec<ParsedSlashCommand>,
 664    invoked_slash_commands: HashMap<InvokedSlashCommandId, InvokedSlashCommand>,
 665    edits_since_last_parse: language::Subscription,
 666    slash_commands: Arc<SlashCommandWorkingSet>,
 667    slash_command_output_sections: Vec<SlashCommandOutputSection<language::Anchor>>,
 668    thought_process_output_sections: Vec<ThoughtProcessOutputSection<language::Anchor>>,
 669    message_anchors: Vec<MessageAnchor>,
 670    contents: Vec<Content>,
 671    messages_metadata: HashMap<MessageId, MessageMetadata>,
 672    summary: ContextSummary,
 673    summary_task: 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    pending_cache_warming_task: Task<Option<()>>,
 680    path: Option<Arc<Path>>,
 681    _subscriptions: Vec<Subscription>,
 682    telemetry: Option<Arc<Telemetry>>,
 683    language_registry: Arc<LanguageRegistry>,
 684    project: Option<Entity<Project>>,
 685    prompt_builder: Arc<PromptBuilder>,
 686    completion_mode: agent_settings::CompletionMode,
 687}
 688
 689trait ContextAnnotation {
 690    fn range(&self) -> &Range<language::Anchor>;
 691}
 692
 693impl ContextAnnotation for ParsedSlashCommand {
 694    fn range(&self) -> &Range<language::Anchor> {
 695        &self.source_range
 696    }
 697}
 698
 699impl EventEmitter<ContextEvent> for AssistantContext {}
 700
 701impl AssistantContext {
 702    pub fn local(
 703        language_registry: Arc<LanguageRegistry>,
 704        project: Option<Entity<Project>>,
 705        telemetry: Option<Arc<Telemetry>>,
 706        prompt_builder: Arc<PromptBuilder>,
 707        slash_commands: Arc<SlashCommandWorkingSet>,
 708        cx: &mut Context<Self>,
 709    ) -> Self {
 710        Self::new(
 711            ContextId::new(),
 712            ReplicaId::default(),
 713            language::Capability::ReadWrite,
 714            language_registry,
 715            prompt_builder,
 716            slash_commands,
 717            project,
 718            telemetry,
 719            cx,
 720        )
 721    }
 722
 723    pub fn completion_mode(&self) -> agent_settings::CompletionMode {
 724        self.completion_mode
 725    }
 726
 727    pub fn set_completion_mode(&mut self, completion_mode: agent_settings::CompletionMode) {
 728        self.completion_mode = completion_mode;
 729    }
 730
 731    pub fn new(
 732        id: ContextId,
 733        replica_id: ReplicaId,
 734        capability: language::Capability,
 735        language_registry: Arc<LanguageRegistry>,
 736        prompt_builder: Arc<PromptBuilder>,
 737        slash_commands: Arc<SlashCommandWorkingSet>,
 738        project: Option<Entity<Project>>,
 739        telemetry: Option<Arc<Telemetry>>,
 740        cx: &mut Context<Self>,
 741    ) -> Self {
 742        let buffer = cx.new(|_cx| {
 743            let buffer = Buffer::remote(
 744                language::BufferId::new(1).unwrap(),
 745                replica_id,
 746                capability,
 747                "",
 748            );
 749            buffer.set_language_registry(language_registry.clone());
 750            buffer
 751        });
 752        let edits_since_last_slash_command_parse =
 753            buffer.update(cx, |buffer, _| buffer.subscribe());
 754        let mut this = Self {
 755            id,
 756            timestamp: clock::Lamport::new(replica_id),
 757            version: clock::Global::new(),
 758            pending_ops: Vec::new(),
 759            operations: Vec::new(),
 760            message_anchors: Default::default(),
 761            contents: Default::default(),
 762            messages_metadata: Default::default(),
 763            parsed_slash_commands: Vec::new(),
 764            invoked_slash_commands: HashMap::default(),
 765            slash_command_output_sections: Vec::new(),
 766            thought_process_output_sections: Vec::new(),
 767            edits_since_last_parse: edits_since_last_slash_command_parse,
 768            summary: ContextSummary::Pending,
 769            summary_task: Task::ready(None),
 770            completion_count: Default::default(),
 771            pending_completions: Default::default(),
 772            token_count: None,
 773            pending_token_count: Task::ready(None),
 774            pending_cache_warming_task: Task::ready(None),
 775            _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
 776            pending_save: Task::ready(Ok(())),
 777            completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,
 778            path: None,
 779            buffer,
 780            telemetry,
 781            project,
 782            language_registry,
 783            slash_commands,
 784            prompt_builder,
 785        };
 786
 787        let first_message_id = MessageId(clock::Lamport {
 788            replica_id: 0,
 789            value: 0,
 790        });
 791        let message = MessageAnchor {
 792            id: first_message_id,
 793            start: language::Anchor::MIN,
 794        };
 795        this.messages_metadata.insert(
 796            first_message_id,
 797            MessageMetadata {
 798                role: Role::User,
 799                status: MessageStatus::Done,
 800                timestamp: first_message_id.0,
 801                cache: None,
 802            },
 803        );
 804        this.message_anchors.push(message);
 805
 806        this.set_language(cx);
 807        this.count_remaining_tokens(cx);
 808        this
 809    }
 810
 811    pub(crate) fn serialize(&self, cx: &App) -> SavedContext {
 812        let buffer = self.buffer.read(cx);
 813        SavedContext {
 814            id: Some(self.id.clone()),
 815            zed: "context".into(),
 816            version: SavedContext::VERSION.into(),
 817            text: buffer.text(),
 818            messages: self
 819                .messages(cx)
 820                .map(|message| SavedMessage {
 821                    id: message.id,
 822                    start: message.offset_range.start,
 823                    metadata: self.messages_metadata[&message.id].clone(),
 824                })
 825                .collect(),
 826            summary: self
 827                .summary
 828                .content()
 829                .map(|summary| summary.text.clone())
 830                .unwrap_or_default(),
 831            slash_command_output_sections: self
 832                .slash_command_output_sections
 833                .iter()
 834                .filter_map(|section| {
 835                    if section.is_valid(buffer) {
 836                        let range = section.range.to_offset(buffer);
 837                        Some(assistant_slash_command::SlashCommandOutputSection {
 838                            range,
 839                            icon: section.icon,
 840                            label: section.label.clone(),
 841                            metadata: section.metadata.clone(),
 842                        })
 843                    } else {
 844                        None
 845                    }
 846                })
 847                .collect(),
 848            thought_process_output_sections: self
 849                .thought_process_output_sections
 850                .iter()
 851                .filter_map(|section| {
 852                    if section.is_valid(buffer) {
 853                        let range = section.range.to_offset(buffer);
 854                        Some(ThoughtProcessOutputSection { range })
 855                    } else {
 856                        None
 857                    }
 858                })
 859                .collect(),
 860        }
 861    }
 862
 863    pub fn deserialize(
 864        saved_context: SavedContext,
 865        path: Arc<Path>,
 866        language_registry: Arc<LanguageRegistry>,
 867        prompt_builder: Arc<PromptBuilder>,
 868        slash_commands: Arc<SlashCommandWorkingSet>,
 869        project: Option<Entity<Project>>,
 870        telemetry: Option<Arc<Telemetry>>,
 871        cx: &mut Context<Self>,
 872    ) -> Self {
 873        let id = saved_context.id.clone().unwrap_or_else(ContextId::new);
 874        let mut this = Self::new(
 875            id,
 876            ReplicaId::default(),
 877            language::Capability::ReadWrite,
 878            language_registry,
 879            prompt_builder,
 880            slash_commands,
 881            project,
 882            telemetry,
 883            cx,
 884        );
 885        this.path = Some(path);
 886        this.buffer.update(cx, |buffer, cx| {
 887            buffer.set_text(saved_context.text.as_str(), cx)
 888        });
 889        let operations = saved_context.into_ops(&this.buffer, cx);
 890        this.apply_ops(operations, cx);
 891        this
 892    }
 893
 894    pub fn id(&self) -> &ContextId {
 895        &self.id
 896    }
 897
 898    pub fn replica_id(&self) -> ReplicaId {
 899        self.timestamp.replica_id
 900    }
 901
 902    pub fn version(&self, cx: &App) -> ContextVersion {
 903        ContextVersion {
 904            context: self.version.clone(),
 905            buffer: self.buffer.read(cx).version(),
 906        }
 907    }
 908
 909    pub fn slash_commands(&self) -> &Arc<SlashCommandWorkingSet> {
 910        &self.slash_commands
 911    }
 912
 913    pub fn set_capability(&mut self, capability: language::Capability, cx: &mut Context<Self>) {
 914        self.buffer
 915            .update(cx, |buffer, cx| buffer.set_capability(capability, cx));
 916    }
 917
 918    fn next_timestamp(&mut self) -> clock::Lamport {
 919        let timestamp = self.timestamp.tick();
 920        self.version.observe(timestamp);
 921        timestamp
 922    }
 923
 924    pub fn serialize_ops(
 925        &self,
 926        since: &ContextVersion,
 927        cx: &App,
 928    ) -> Task<Vec<proto::ContextOperation>> {
 929        let buffer_ops = self
 930            .buffer
 931            .read(cx)
 932            .serialize_ops(Some(since.buffer.clone()), cx);
 933
 934        let mut context_ops = self
 935            .operations
 936            .iter()
 937            .filter(|op| !since.context.observed(op.timestamp()))
 938            .cloned()
 939            .collect::<Vec<_>>();
 940        context_ops.extend(self.pending_ops.iter().cloned());
 941
 942        cx.background_spawn(async move {
 943            let buffer_ops = buffer_ops.await;
 944            context_ops.sort_unstable_by_key(|op| op.timestamp());
 945            buffer_ops
 946                .into_iter()
 947                .map(|op| proto::ContextOperation {
 948                    variant: Some(proto::context_operation::Variant::BufferOperation(
 949                        proto::context_operation::BufferOperation {
 950                            operation: Some(op),
 951                        },
 952                    )),
 953                })
 954                .chain(context_ops.into_iter().map(|op| op.to_proto()))
 955                .collect()
 956        })
 957    }
 958
 959    pub fn apply_ops(
 960        &mut self,
 961        ops: impl IntoIterator<Item = ContextOperation>,
 962        cx: &mut Context<Self>,
 963    ) {
 964        let mut buffer_ops = Vec::new();
 965        for op in ops {
 966            match op {
 967                ContextOperation::BufferOperation(buffer_op) => buffer_ops.push(buffer_op),
 968                op @ _ => self.pending_ops.push(op),
 969            }
 970        }
 971        self.buffer
 972            .update(cx, |buffer, cx| buffer.apply_ops(buffer_ops, cx));
 973        self.flush_ops(cx);
 974    }
 975
 976    fn flush_ops(&mut self, cx: &mut Context<AssistantContext>) {
 977        let mut changed_messages = HashSet::default();
 978        let mut summary_generated = false;
 979
 980        self.pending_ops.sort_unstable_by_key(|op| op.timestamp());
 981        for op in mem::take(&mut self.pending_ops) {
 982            if !self.can_apply_op(&op, cx) {
 983                self.pending_ops.push(op);
 984                continue;
 985            }
 986
 987            let timestamp = op.timestamp();
 988            match op.clone() {
 989                ContextOperation::InsertMessage {
 990                    anchor, metadata, ..
 991                } => {
 992                    if self.messages_metadata.contains_key(&anchor.id) {
 993                        // We already applied this operation.
 994                    } else {
 995                        changed_messages.insert(anchor.id);
 996                        self.insert_message(anchor, metadata, cx);
 997                    }
 998                }
 999                ContextOperation::UpdateMessage {
1000                    message_id,
1001                    metadata: new_metadata,
1002                    ..
1003                } => {
1004                    let metadata = self.messages_metadata.get_mut(&message_id).unwrap();
1005                    if new_metadata.timestamp > metadata.timestamp {
1006                        *metadata = new_metadata;
1007                        changed_messages.insert(message_id);
1008                    }
1009                }
1010                ContextOperation::UpdateSummary {
1011                    summary: new_summary,
1012                    ..
1013                } => {
1014                    if self.summary.timestamp().map_or(true, |current_timestamp| {
1015                        new_summary.timestamp > current_timestamp
1016                    }) {
1017                        self.summary = ContextSummary::Content(new_summary);
1018                        summary_generated = true;
1019                    }
1020                }
1021                ContextOperation::SlashCommandStarted {
1022                    id,
1023                    output_range,
1024                    name,
1025                    ..
1026                } => {
1027                    self.invoked_slash_commands.insert(
1028                        id,
1029                        InvokedSlashCommand {
1030                            name: name.into(),
1031                            range: output_range,
1032                            run_commands_in_ranges: Vec::new(),
1033                            status: InvokedSlashCommandStatus::Running(Task::ready(())),
1034                            transaction: None,
1035                            timestamp: id.0,
1036                        },
1037                    );
1038                    cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id: id });
1039                }
1040                ContextOperation::SlashCommandOutputSectionAdded { section, .. } => {
1041                    let buffer = self.buffer.read(cx);
1042                    if let Err(ix) = self
1043                        .slash_command_output_sections
1044                        .binary_search_by(|probe| probe.range.cmp(&section.range, buffer))
1045                    {
1046                        self.slash_command_output_sections
1047                            .insert(ix, section.clone());
1048                        cx.emit(ContextEvent::SlashCommandOutputSectionAdded { section });
1049                    }
1050                }
1051                ContextOperation::ThoughtProcessOutputSectionAdded { section, .. } => {
1052                    let buffer = self.buffer.read(cx);
1053                    if let Err(ix) = self
1054                        .thought_process_output_sections
1055                        .binary_search_by(|probe| probe.range.cmp(&section.range, buffer))
1056                    {
1057                        self.thought_process_output_sections
1058                            .insert(ix, section.clone());
1059                    }
1060                }
1061                ContextOperation::SlashCommandFinished {
1062                    id,
1063                    error_message,
1064                    timestamp,
1065                    ..
1066                } => {
1067                    if let Some(slash_command) = self.invoked_slash_commands.get_mut(&id) {
1068                        if timestamp > slash_command.timestamp {
1069                            slash_command.timestamp = timestamp;
1070                            match error_message {
1071                                Some(message) => {
1072                                    slash_command.status =
1073                                        InvokedSlashCommandStatus::Error(message.into());
1074                                }
1075                                None => {
1076                                    slash_command.status = InvokedSlashCommandStatus::Finished;
1077                                }
1078                            }
1079                            cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id: id });
1080                        }
1081                    }
1082                }
1083                ContextOperation::BufferOperation(_) => unreachable!(),
1084            }
1085
1086            self.version.observe(timestamp);
1087            self.timestamp.observe(timestamp);
1088            self.operations.push(op);
1089        }
1090
1091        if !changed_messages.is_empty() {
1092            self.message_roles_updated(changed_messages, cx);
1093            cx.emit(ContextEvent::MessagesEdited);
1094            cx.notify();
1095        }
1096
1097        if summary_generated {
1098            cx.emit(ContextEvent::SummaryChanged);
1099            cx.emit(ContextEvent::SummaryGenerated);
1100            cx.notify();
1101        }
1102    }
1103
1104    fn can_apply_op(&self, op: &ContextOperation, cx: &App) -> bool {
1105        if !self.version.observed_all(op.version()) {
1106            return false;
1107        }
1108
1109        match op {
1110            ContextOperation::InsertMessage { anchor, .. } => self
1111                .buffer
1112                .read(cx)
1113                .version
1114                .observed(anchor.start.timestamp),
1115            ContextOperation::UpdateMessage { message_id, .. } => {
1116                self.messages_metadata.contains_key(message_id)
1117            }
1118            ContextOperation::UpdateSummary { .. } => true,
1119            ContextOperation::SlashCommandStarted { output_range, .. } => {
1120                self.has_received_operations_for_anchor_range(output_range.clone(), cx)
1121            }
1122            ContextOperation::SlashCommandOutputSectionAdded { section, .. } => {
1123                self.has_received_operations_for_anchor_range(section.range.clone(), cx)
1124            }
1125            ContextOperation::ThoughtProcessOutputSectionAdded { section, .. } => {
1126                self.has_received_operations_for_anchor_range(section.range.clone(), cx)
1127            }
1128            ContextOperation::SlashCommandFinished { .. } => true,
1129            ContextOperation::BufferOperation(_) => {
1130                panic!("buffer operations should always be applied")
1131            }
1132        }
1133    }
1134
1135    fn has_received_operations_for_anchor_range(
1136        &self,
1137        range: Range<text::Anchor>,
1138        cx: &App,
1139    ) -> bool {
1140        let version = &self.buffer.read(cx).version;
1141        let observed_start = range.start == language::Anchor::MIN
1142            || range.start == language::Anchor::MAX
1143            || version.observed(range.start.timestamp);
1144        let observed_end = range.end == language::Anchor::MIN
1145            || range.end == language::Anchor::MAX
1146            || version.observed(range.end.timestamp);
1147        observed_start && observed_end
1148    }
1149
1150    fn push_op(&mut self, op: ContextOperation, cx: &mut Context<Self>) {
1151        self.operations.push(op.clone());
1152        cx.emit(ContextEvent::Operation(op));
1153    }
1154
1155    pub fn buffer(&self) -> &Entity<Buffer> {
1156        &self.buffer
1157    }
1158
1159    pub fn language_registry(&self) -> Arc<LanguageRegistry> {
1160        self.language_registry.clone()
1161    }
1162
1163    pub fn project(&self) -> Option<Entity<Project>> {
1164        self.project.clone()
1165    }
1166
1167    pub fn prompt_builder(&self) -> Arc<PromptBuilder> {
1168        self.prompt_builder.clone()
1169    }
1170
1171    pub fn path(&self) -> Option<&Arc<Path>> {
1172        self.path.as_ref()
1173    }
1174
1175    pub fn summary(&self) -> &ContextSummary {
1176        &self.summary
1177    }
1178
1179    pub fn parsed_slash_commands(&self) -> &[ParsedSlashCommand] {
1180        &self.parsed_slash_commands
1181    }
1182
1183    pub fn invoked_slash_command(
1184        &self,
1185        command_id: &InvokedSlashCommandId,
1186    ) -> Option<&InvokedSlashCommand> {
1187        self.invoked_slash_commands.get(command_id)
1188    }
1189
1190    pub fn slash_command_output_sections(&self) -> &[SlashCommandOutputSection<language::Anchor>] {
1191        &self.slash_command_output_sections
1192    }
1193
1194    pub fn thought_process_output_sections(
1195        &self,
1196    ) -> &[ThoughtProcessOutputSection<language::Anchor>] {
1197        &self.thought_process_output_sections
1198    }
1199
1200    pub fn contains_files(&self, cx: &App) -> bool {
1201        let buffer = self.buffer.read(cx);
1202        self.slash_command_output_sections.iter().any(|section| {
1203            section.is_valid(buffer)
1204                && section
1205                    .metadata
1206                    .as_ref()
1207                    .and_then(|metadata| {
1208                        serde_json::from_value::<FileCommandMetadata>(metadata.clone()).ok()
1209                    })
1210                    .is_some()
1211        })
1212    }
1213
1214    fn set_language(&mut self, cx: &mut Context<Self>) {
1215        let markdown = self.language_registry.language_for_name("Markdown");
1216        cx.spawn(async move |this, cx| {
1217            let markdown = markdown.await?;
1218            this.update(cx, |this, cx| {
1219                this.buffer
1220                    .update(cx, |buffer, cx| buffer.set_language(Some(markdown), cx));
1221            })
1222        })
1223        .detach_and_log_err(cx);
1224    }
1225
1226    fn handle_buffer_event(
1227        &mut self,
1228        _: Entity<Buffer>,
1229        event: &language::BufferEvent,
1230        cx: &mut Context<Self>,
1231    ) {
1232        match event {
1233            language::BufferEvent::Operation {
1234                operation,
1235                is_local: true,
1236            } => cx.emit(ContextEvent::Operation(ContextOperation::BufferOperation(
1237                operation.clone(),
1238            ))),
1239            language::BufferEvent::Edited => {
1240                self.count_remaining_tokens(cx);
1241                self.reparse(cx);
1242                cx.emit(ContextEvent::MessagesEdited);
1243            }
1244            _ => {}
1245        }
1246    }
1247
1248    pub fn token_count(&self) -> Option<usize> {
1249        self.token_count
1250    }
1251
1252    pub(crate) fn count_remaining_tokens(&mut self, cx: &mut Context<Self>) {
1253        // Assume it will be a Chat request, even though that takes fewer tokens (and risks going over the limit),
1254        // because otherwise you see in the UI that your empty message has a bunch of tokens already used.
1255        let Some(model) = LanguageModelRegistry::read_global(cx).default_model() else {
1256            return;
1257        };
1258        let request = self.to_completion_request(Some(&model.model), cx);
1259        let debounce = self.token_count.is_some();
1260        self.pending_token_count = cx.spawn(async move |this, cx| {
1261            async move {
1262                if debounce {
1263                    cx.background_executor()
1264                        .timer(Duration::from_millis(200))
1265                        .await;
1266                }
1267
1268                let token_count = cx
1269                    .update(|cx| model.model.count_tokens(request, cx))?
1270                    .await?;
1271                this.update(cx, |this, cx| {
1272                    this.token_count = Some(token_count);
1273                    this.start_cache_warming(&model.model, cx);
1274                    cx.notify()
1275                })
1276            }
1277            .log_err()
1278            .await
1279        });
1280    }
1281
1282    pub fn mark_cache_anchors(
1283        &mut self,
1284        cache_configuration: &Option<LanguageModelCacheConfiguration>,
1285        speculative: bool,
1286        cx: &mut Context<Self>,
1287    ) -> bool {
1288        let cache_configuration =
1289            cache_configuration
1290                .as_ref()
1291                .unwrap_or(&LanguageModelCacheConfiguration {
1292                    max_cache_anchors: 0,
1293                    should_speculate: false,
1294                    min_total_token: 0,
1295                });
1296
1297        let messages: Vec<Message> = self.messages(cx).collect();
1298
1299        let mut sorted_messages = messages.clone();
1300        if speculative {
1301            // Avoid caching the last message if this is a speculative cache fetch as
1302            // it's likely to change.
1303            sorted_messages.pop();
1304        }
1305        sorted_messages.retain(|m| m.role == Role::User);
1306        sorted_messages.sort_by(|a, b| b.offset_range.len().cmp(&a.offset_range.len()));
1307
1308        let cache_anchors = if self.token_count.unwrap_or(0) < cache_configuration.min_total_token {
1309            // If we have't hit the minimum threshold to enable caching, don't cache anything.
1310            0
1311        } else {
1312            // Save 1 anchor for the inline assistant to use.
1313            max(cache_configuration.max_cache_anchors, 1) - 1
1314        };
1315        sorted_messages.truncate(cache_anchors);
1316
1317        let anchors: HashSet<MessageId> = sorted_messages
1318            .into_iter()
1319            .map(|message| message.id)
1320            .collect();
1321
1322        let buffer = self.buffer.read(cx).snapshot();
1323        let invalidated_caches: HashSet<MessageId> = messages
1324            .iter()
1325            .scan(false, |encountered_invalid, message| {
1326                let message_id = message.id;
1327                let is_invalid = self
1328                    .messages_metadata
1329                    .get(&message_id)
1330                    .map_or(true, |metadata| {
1331                        !metadata.is_cache_valid(&buffer, &message.offset_range)
1332                            || *encountered_invalid
1333                    });
1334                *encountered_invalid |= is_invalid;
1335                Some(if is_invalid { Some(message_id) } else { None })
1336            })
1337            .flatten()
1338            .collect();
1339
1340        let last_anchor = messages.iter().rev().find_map(|message| {
1341            if anchors.contains(&message.id) {
1342                Some(message.id)
1343            } else {
1344                None
1345            }
1346        });
1347
1348        let mut new_anchor_needs_caching = false;
1349        let current_version = &buffer.version;
1350        // If we have no anchors, mark all messages as not being cached.
1351        let mut hit_last_anchor = last_anchor.is_none();
1352
1353        for message in messages.iter() {
1354            if hit_last_anchor {
1355                self.update_metadata(message.id, cx, |metadata| metadata.cache = None);
1356                continue;
1357            }
1358
1359            if let Some(last_anchor) = last_anchor {
1360                if message.id == last_anchor {
1361                    hit_last_anchor = true;
1362                }
1363            }
1364
1365            new_anchor_needs_caching = new_anchor_needs_caching
1366                || (invalidated_caches.contains(&message.id) && anchors.contains(&message.id));
1367
1368            self.update_metadata(message.id, cx, |metadata| {
1369                let cache_status = if invalidated_caches.contains(&message.id) {
1370                    CacheStatus::Pending
1371                } else {
1372                    metadata
1373                        .cache
1374                        .as_ref()
1375                        .map_or(CacheStatus::Pending, |cm| cm.status.clone())
1376                };
1377                metadata.cache = Some(MessageCacheMetadata {
1378                    is_anchor: anchors.contains(&message.id),
1379                    is_final_anchor: hit_last_anchor,
1380                    status: cache_status,
1381                    cached_at: current_version.clone(),
1382                });
1383            });
1384        }
1385        new_anchor_needs_caching
1386    }
1387
1388    fn start_cache_warming(&mut self, model: &Arc<dyn LanguageModel>, cx: &mut Context<Self>) {
1389        let cache_configuration = model.cache_configuration();
1390
1391        if !self.mark_cache_anchors(&cache_configuration, true, cx) {
1392            return;
1393        }
1394        if !self.pending_completions.is_empty() {
1395            return;
1396        }
1397        if let Some(cache_configuration) = cache_configuration {
1398            if !cache_configuration.should_speculate {
1399                return;
1400            }
1401        }
1402
1403        let request = {
1404            let mut req = self.to_completion_request(Some(&model), cx);
1405            // Skip the last message because it's likely to change and
1406            // therefore would be a waste to cache.
1407            req.messages.pop();
1408            req.messages.push(LanguageModelRequestMessage {
1409                role: Role::User,
1410                content: vec!["Respond only with OK, nothing else.".into()],
1411                cache: false,
1412            });
1413            req
1414        };
1415
1416        let model = Arc::clone(model);
1417        self.pending_cache_warming_task = cx.spawn(async move |this, cx| {
1418            async move {
1419                match model.stream_completion(request, &cx).await {
1420                    Ok(mut stream) => {
1421                        stream.next().await;
1422                        log::info!("Cache warming completed successfully");
1423                    }
1424                    Err(e) => {
1425                        log::warn!("Cache warming failed: {}", e);
1426                    }
1427                };
1428                this.update(cx, |this, cx| {
1429                    this.update_cache_status_for_completion(cx);
1430                })
1431                .ok();
1432                anyhow::Ok(())
1433            }
1434            .log_err()
1435            .await
1436        });
1437    }
1438
1439    pub fn update_cache_status_for_completion(&mut self, cx: &mut Context<Self>) {
1440        let cached_message_ids: Vec<MessageId> = self
1441            .messages_metadata
1442            .iter()
1443            .filter_map(|(message_id, metadata)| {
1444                metadata.cache.as_ref().and_then(|cache| {
1445                    if cache.status == CacheStatus::Pending {
1446                        Some(*message_id)
1447                    } else {
1448                        None
1449                    }
1450                })
1451            })
1452            .collect();
1453
1454        for message_id in cached_message_ids {
1455            self.update_metadata(message_id, cx, |metadata| {
1456                if let Some(cache) = &mut metadata.cache {
1457                    cache.status = CacheStatus::Cached;
1458                }
1459            });
1460        }
1461        cx.notify();
1462    }
1463
1464    pub fn reparse(&mut self, cx: &mut Context<Self>) {
1465        let buffer = self.buffer.read(cx).text_snapshot();
1466        let mut row_ranges = self
1467            .edits_since_last_parse
1468            .consume()
1469            .into_iter()
1470            .map(|edit| {
1471                let start_row = buffer.offset_to_point(edit.new.start).row;
1472                let end_row = buffer.offset_to_point(edit.new.end).row + 1;
1473                start_row..end_row
1474            })
1475            .peekable();
1476
1477        let mut removed_parsed_slash_command_ranges = Vec::new();
1478        let mut updated_parsed_slash_commands = Vec::new();
1479        while let Some(mut row_range) = row_ranges.next() {
1480            while let Some(next_row_range) = row_ranges.peek() {
1481                if row_range.end >= next_row_range.start {
1482                    row_range.end = next_row_range.end;
1483                    row_ranges.next();
1484                } else {
1485                    break;
1486                }
1487            }
1488
1489            let start = buffer.anchor_before(Point::new(row_range.start, 0));
1490            let end = buffer.anchor_after(Point::new(
1491                row_range.end - 1,
1492                buffer.line_len(row_range.end - 1),
1493            ));
1494
1495            self.reparse_slash_commands_in_range(
1496                start..end,
1497                &buffer,
1498                &mut updated_parsed_slash_commands,
1499                &mut removed_parsed_slash_command_ranges,
1500                cx,
1501            );
1502            self.invalidate_pending_slash_commands(&buffer, cx);
1503        }
1504
1505        if !updated_parsed_slash_commands.is_empty()
1506            || !removed_parsed_slash_command_ranges.is_empty()
1507        {
1508            cx.emit(ContextEvent::ParsedSlashCommandsUpdated {
1509                removed: removed_parsed_slash_command_ranges,
1510                updated: updated_parsed_slash_commands,
1511            });
1512        }
1513    }
1514
1515    fn reparse_slash_commands_in_range(
1516        &mut self,
1517        range: Range<text::Anchor>,
1518        buffer: &BufferSnapshot,
1519        updated: &mut Vec<ParsedSlashCommand>,
1520        removed: &mut Vec<Range<text::Anchor>>,
1521        cx: &App,
1522    ) {
1523        let old_range = self.pending_command_indices_for_range(range.clone(), cx);
1524
1525        let mut new_commands = Vec::new();
1526        let mut lines = buffer.text_for_range(range).lines();
1527        let mut offset = lines.offset();
1528        while let Some(line) = lines.next() {
1529            if let Some(command_line) = SlashCommandLine::parse(line) {
1530                let name = &line[command_line.name.clone()];
1531                let arguments = command_line
1532                    .arguments
1533                    .iter()
1534                    .filter_map(|argument_range| {
1535                        if argument_range.is_empty() {
1536                            None
1537                        } else {
1538                            line.get(argument_range.clone())
1539                        }
1540                    })
1541                    .map(ToOwned::to_owned)
1542                    .collect::<SmallVec<_>>();
1543                if let Some(command) = self.slash_commands.command(name, cx) {
1544                    if !command.requires_argument() || !arguments.is_empty() {
1545                        let start_ix = offset + command_line.name.start - 1;
1546                        let end_ix = offset
1547                            + command_line
1548                                .arguments
1549                                .last()
1550                                .map_or(command_line.name.end, |argument| argument.end);
1551                        let source_range =
1552                            buffer.anchor_after(start_ix)..buffer.anchor_after(end_ix);
1553                        let pending_command = ParsedSlashCommand {
1554                            name: name.to_string(),
1555                            arguments,
1556                            source_range,
1557                            status: PendingSlashCommandStatus::Idle,
1558                        };
1559                        updated.push(pending_command.clone());
1560                        new_commands.push(pending_command);
1561                    }
1562                }
1563            }
1564
1565            offset = lines.offset();
1566        }
1567
1568        let removed_commands = self.parsed_slash_commands.splice(old_range, new_commands);
1569        removed.extend(removed_commands.map(|command| command.source_range));
1570    }
1571
1572    fn invalidate_pending_slash_commands(
1573        &mut self,
1574        buffer: &BufferSnapshot,
1575        cx: &mut Context<Self>,
1576    ) {
1577        let mut invalidated_command_ids = Vec::new();
1578        for (&command_id, command) in self.invoked_slash_commands.iter_mut() {
1579            if !matches!(command.status, InvokedSlashCommandStatus::Finished)
1580                && (!command.range.start.is_valid(buffer) || !command.range.end.is_valid(buffer))
1581            {
1582                command.status = InvokedSlashCommandStatus::Finished;
1583                cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id });
1584                invalidated_command_ids.push(command_id);
1585            }
1586        }
1587
1588        for command_id in invalidated_command_ids {
1589            let version = self.version.clone();
1590            let timestamp = self.next_timestamp();
1591            self.push_op(
1592                ContextOperation::SlashCommandFinished {
1593                    id: command_id,
1594                    timestamp,
1595                    error_message: None,
1596                    version: version.clone(),
1597                },
1598                cx,
1599            );
1600        }
1601    }
1602
1603    pub fn pending_command_for_position(
1604        &mut self,
1605        position: language::Anchor,
1606        cx: &mut Context<Self>,
1607    ) -> Option<&mut ParsedSlashCommand> {
1608        let buffer = self.buffer.read(cx);
1609        match self
1610            .parsed_slash_commands
1611            .binary_search_by(|probe| probe.source_range.end.cmp(&position, buffer))
1612        {
1613            Ok(ix) => Some(&mut self.parsed_slash_commands[ix]),
1614            Err(ix) => {
1615                let cmd = self.parsed_slash_commands.get_mut(ix)?;
1616                if position.cmp(&cmd.source_range.start, buffer).is_ge()
1617                    && position.cmp(&cmd.source_range.end, buffer).is_le()
1618                {
1619                    Some(cmd)
1620                } else {
1621                    None
1622                }
1623            }
1624        }
1625    }
1626
1627    pub fn pending_commands_for_range(
1628        &self,
1629        range: Range<language::Anchor>,
1630        cx: &App,
1631    ) -> &[ParsedSlashCommand] {
1632        let range = self.pending_command_indices_for_range(range, cx);
1633        &self.parsed_slash_commands[range]
1634    }
1635
1636    fn pending_command_indices_for_range(
1637        &self,
1638        range: Range<language::Anchor>,
1639        cx: &App,
1640    ) -> Range<usize> {
1641        self.indices_intersecting_buffer_range(&self.parsed_slash_commands, range, cx)
1642    }
1643
1644    fn indices_intersecting_buffer_range<T: ContextAnnotation>(
1645        &self,
1646        all_annotations: &[T],
1647        range: Range<language::Anchor>,
1648        cx: &App,
1649    ) -> Range<usize> {
1650        let buffer = self.buffer.read(cx);
1651        let start_ix = match all_annotations
1652            .binary_search_by(|probe| probe.range().end.cmp(&range.start, &buffer))
1653        {
1654            Ok(ix) | Err(ix) => ix,
1655        };
1656        let end_ix = match all_annotations
1657            .binary_search_by(|probe| probe.range().start.cmp(&range.end, &buffer))
1658        {
1659            Ok(ix) => ix + 1,
1660            Err(ix) => ix,
1661        };
1662        start_ix..end_ix
1663    }
1664
1665    pub fn insert_command_output(
1666        &mut self,
1667        command_source_range: Range<language::Anchor>,
1668        name: &str,
1669        output: Task<SlashCommandResult>,
1670        ensure_trailing_newline: bool,
1671        cx: &mut Context<Self>,
1672    ) {
1673        let version = self.version.clone();
1674        let command_id = InvokedSlashCommandId(self.next_timestamp());
1675
1676        const PENDING_OUTPUT_END_MARKER: &str = "";
1677
1678        let (command_range, command_source_range, insert_position, first_transaction) =
1679            self.buffer.update(cx, |buffer, cx| {
1680                let command_source_range = command_source_range.to_offset(buffer);
1681                let mut insertion = format!("\n{PENDING_OUTPUT_END_MARKER}");
1682                if ensure_trailing_newline {
1683                    insertion.push('\n');
1684                }
1685
1686                buffer.finalize_last_transaction();
1687                buffer.start_transaction();
1688                buffer.edit(
1689                    [(
1690                        command_source_range.end..command_source_range.end,
1691                        insertion,
1692                    )],
1693                    None,
1694                    cx,
1695                );
1696                let first_transaction = buffer.end_transaction(cx).unwrap();
1697                buffer.finalize_last_transaction();
1698
1699                let insert_position = buffer.anchor_after(command_source_range.end + 1);
1700                let command_range = buffer.anchor_after(command_source_range.start)
1701                    ..buffer.anchor_before(
1702                        command_source_range.end + 1 + PENDING_OUTPUT_END_MARKER.len(),
1703                    );
1704                let command_source_range = buffer.anchor_before(command_source_range.start)
1705                    ..buffer.anchor_before(command_source_range.end + 1);
1706                (
1707                    command_range,
1708                    command_source_range,
1709                    insert_position,
1710                    first_transaction,
1711                )
1712            });
1713        self.reparse(cx);
1714
1715        let insert_output_task = cx.spawn(async move |this, cx| {
1716            let run_command = async {
1717                let mut stream = output.await?;
1718
1719                struct PendingSection {
1720                    start: language::Anchor,
1721                    icon: IconName,
1722                    label: SharedString,
1723                    metadata: Option<serde_json::Value>,
1724                }
1725
1726                let mut pending_section_stack: Vec<PendingSection> = Vec::new();
1727                let mut last_role: Option<Role> = None;
1728                let mut last_section_range = None;
1729
1730                while let Some(event) = stream.next().await {
1731                    let event = event?;
1732                    this.update(cx, |this, cx| {
1733                        this.buffer.update(cx, |buffer, _cx| {
1734                            buffer.finalize_last_transaction();
1735                            buffer.start_transaction()
1736                        });
1737
1738                        match event {
1739                            SlashCommandEvent::StartMessage {
1740                                role,
1741                                merge_same_roles,
1742                            } => {
1743                                if !merge_same_roles && Some(role) != last_role {
1744                                    let buffer = this.buffer.read(cx);
1745                                    let offset = insert_position.to_offset(buffer);
1746                                    this.insert_message_at_offset(
1747                                        offset,
1748                                        role,
1749                                        MessageStatus::Pending,
1750                                        cx,
1751                                    );
1752                                }
1753
1754                                last_role = Some(role);
1755                            }
1756                            SlashCommandEvent::StartSection {
1757                                icon,
1758                                label,
1759                                metadata,
1760                            } => {
1761                                this.buffer.update(cx, |buffer, cx| {
1762                                    let insert_point = insert_position.to_point(buffer);
1763                                    if insert_point.column > 0 {
1764                                        buffer.edit([(insert_point..insert_point, "\n")], None, cx);
1765                                    }
1766
1767                                    pending_section_stack.push(PendingSection {
1768                                        start: buffer.anchor_before(insert_position),
1769                                        icon,
1770                                        label,
1771                                        metadata,
1772                                    });
1773                                });
1774                            }
1775                            SlashCommandEvent::Content(SlashCommandContent::Text {
1776                                text,
1777                                run_commands_in_text,
1778                            }) => {
1779                                let start = this.buffer.read(cx).anchor_before(insert_position);
1780
1781                                this.buffer.update(cx, |buffer, cx| {
1782                                    buffer.edit(
1783                                        [(insert_position..insert_position, text)],
1784                                        None,
1785                                        cx,
1786                                    )
1787                                });
1788
1789                                let end = this.buffer.read(cx).anchor_before(insert_position);
1790                                if run_commands_in_text {
1791                                    if let Some(invoked_slash_command) =
1792                                        this.invoked_slash_commands.get_mut(&command_id)
1793                                    {
1794                                        invoked_slash_command
1795                                            .run_commands_in_ranges
1796                                            .push(start..end);
1797                                    }
1798                                }
1799                            }
1800                            SlashCommandEvent::EndSection => {
1801                                if let Some(pending_section) = pending_section_stack.pop() {
1802                                    let offset_range = (pending_section.start..insert_position)
1803                                        .to_offset(this.buffer.read(cx));
1804                                    if !offset_range.is_empty() {
1805                                        let range = this.buffer.update(cx, |buffer, _cx| {
1806                                            buffer.anchor_after(offset_range.start)
1807                                                ..buffer.anchor_before(offset_range.end)
1808                                        });
1809                                        this.insert_slash_command_output_section(
1810                                            SlashCommandOutputSection {
1811                                                range: range.clone(),
1812                                                icon: pending_section.icon,
1813                                                label: pending_section.label,
1814                                                metadata: pending_section.metadata,
1815                                            },
1816                                            cx,
1817                                        );
1818                                        last_section_range = Some(range);
1819                                    }
1820                                }
1821                            }
1822                        }
1823
1824                        this.buffer.update(cx, |buffer, cx| {
1825                            if let Some(event_transaction) = buffer.end_transaction(cx) {
1826                                buffer.merge_transactions(event_transaction, first_transaction);
1827                            }
1828                        });
1829                    })?;
1830                }
1831
1832                this.update(cx, |this, cx| {
1833                    this.buffer.update(cx, |buffer, cx| {
1834                        buffer.finalize_last_transaction();
1835                        buffer.start_transaction();
1836
1837                        let mut deletions = vec![(command_source_range.to_offset(buffer), "")];
1838                        let insert_position = insert_position.to_offset(buffer);
1839                        let command_range_end = command_range.end.to_offset(buffer);
1840
1841                        if buffer.contains_str_at(insert_position, PENDING_OUTPUT_END_MARKER) {
1842                            deletions.push((
1843                                insert_position..insert_position + PENDING_OUTPUT_END_MARKER.len(),
1844                                "",
1845                            ));
1846                        }
1847
1848                        if ensure_trailing_newline
1849                            && buffer.contains_str_at(command_range_end, "\n")
1850                        {
1851                            let newline_offset = insert_position.saturating_sub(1);
1852                            if buffer.contains_str_at(newline_offset, "\n")
1853                                && last_section_range.map_or(true, |last_section_range| {
1854                                    !last_section_range
1855                                        .to_offset(buffer)
1856                                        .contains(&newline_offset)
1857                                })
1858                            {
1859                                deletions.push((command_range_end..command_range_end + 1, ""));
1860                            }
1861                        }
1862
1863                        buffer.edit(deletions, None, cx);
1864
1865                        if let Some(deletion_transaction) = buffer.end_transaction(cx) {
1866                            buffer.merge_transactions(deletion_transaction, first_transaction);
1867                        }
1868                    });
1869                })?;
1870
1871                debug_assert!(pending_section_stack.is_empty());
1872
1873                anyhow::Ok(())
1874            };
1875
1876            let command_result = run_command.await;
1877
1878            this.update(cx, |this, cx| {
1879                let version = this.version.clone();
1880                let timestamp = this.next_timestamp();
1881                let Some(invoked_slash_command) = this.invoked_slash_commands.get_mut(&command_id)
1882                else {
1883                    return;
1884                };
1885                let mut error_message = None;
1886                match command_result {
1887                    Ok(()) => {
1888                        invoked_slash_command.status = InvokedSlashCommandStatus::Finished;
1889                    }
1890                    Err(error) => {
1891                        let message = error.to_string();
1892                        invoked_slash_command.status =
1893                            InvokedSlashCommandStatus::Error(message.clone().into());
1894                        error_message = Some(message);
1895                    }
1896                }
1897
1898                cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id });
1899                this.push_op(
1900                    ContextOperation::SlashCommandFinished {
1901                        id: command_id,
1902                        timestamp,
1903                        error_message,
1904                        version,
1905                    },
1906                    cx,
1907                );
1908            })
1909            .ok();
1910        });
1911
1912        self.invoked_slash_commands.insert(
1913            command_id,
1914            InvokedSlashCommand {
1915                name: name.to_string().into(),
1916                range: command_range.clone(),
1917                run_commands_in_ranges: Vec::new(),
1918                status: InvokedSlashCommandStatus::Running(insert_output_task),
1919                transaction: Some(first_transaction),
1920                timestamp: command_id.0,
1921            },
1922        );
1923        cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id });
1924        self.push_op(
1925            ContextOperation::SlashCommandStarted {
1926                id: command_id,
1927                output_range: command_range,
1928                name: name.to_string(),
1929                version,
1930            },
1931            cx,
1932        );
1933    }
1934
1935    fn insert_slash_command_output_section(
1936        &mut self,
1937        section: SlashCommandOutputSection<language::Anchor>,
1938        cx: &mut Context<Self>,
1939    ) {
1940        let buffer = self.buffer.read(cx);
1941        let insertion_ix = match self
1942            .slash_command_output_sections
1943            .binary_search_by(|probe| probe.range.cmp(&section.range, buffer))
1944        {
1945            Ok(ix) | Err(ix) => ix,
1946        };
1947        self.slash_command_output_sections
1948            .insert(insertion_ix, section.clone());
1949        cx.emit(ContextEvent::SlashCommandOutputSectionAdded {
1950            section: section.clone(),
1951        });
1952        let version = self.version.clone();
1953        let timestamp = self.next_timestamp();
1954        self.push_op(
1955            ContextOperation::SlashCommandOutputSectionAdded {
1956                timestamp,
1957                section,
1958                version,
1959            },
1960            cx,
1961        );
1962    }
1963
1964    fn insert_thought_process_output_section(
1965        &mut self,
1966        section: ThoughtProcessOutputSection<language::Anchor>,
1967        cx: &mut Context<Self>,
1968    ) {
1969        let buffer = self.buffer.read(cx);
1970        let insertion_ix = match self
1971            .thought_process_output_sections
1972            .binary_search_by(|probe| probe.range.cmp(&section.range, buffer))
1973        {
1974            Ok(ix) | Err(ix) => ix,
1975        };
1976        self.thought_process_output_sections
1977            .insert(insertion_ix, section.clone());
1978        // cx.emit(ContextEvent::ThoughtProcessOutputSectionAdded {
1979        //     section: section.clone(),
1980        // });
1981        let version = self.version.clone();
1982        let timestamp = self.next_timestamp();
1983        self.push_op(
1984            ContextOperation::ThoughtProcessOutputSectionAdded {
1985                timestamp,
1986                section,
1987                version,
1988            },
1989            cx,
1990        );
1991    }
1992
1993    pub fn completion_provider_changed(&mut self, cx: &mut Context<Self>) {
1994        self.count_remaining_tokens(cx);
1995    }
1996
1997    fn get_last_valid_message_id(&self, cx: &Context<Self>) -> Option<MessageId> {
1998        self.message_anchors.iter().rev().find_map(|message| {
1999            message
2000                .start
2001                .is_valid(self.buffer.read(cx))
2002                .then_some(message.id)
2003        })
2004    }
2005
2006    pub fn assist(&mut self, cx: &mut Context<Self>) -> Option<MessageAnchor> {
2007        let model_registry = LanguageModelRegistry::read_global(cx);
2008        let model = model_registry.default_model()?;
2009        let last_message_id = self.get_last_valid_message_id(cx)?;
2010
2011        if !model.provider.is_authenticated(cx) {
2012            log::info!("completion provider has no credentials");
2013            return None;
2014        }
2015
2016        let model = model.model;
2017
2018        // Compute which messages to cache, including the last one.
2019        self.mark_cache_anchors(&model.cache_configuration(), false, cx);
2020
2021        let request = self.to_completion_request(Some(&model), cx);
2022
2023        let assistant_message = self
2024            .insert_message_after(last_message_id, Role::Assistant, MessageStatus::Pending, cx)
2025            .unwrap();
2026
2027        // Queue up the user's next reply.
2028        let user_message = self
2029            .insert_message_after(assistant_message.id, Role::User, MessageStatus::Done, cx)
2030            .unwrap();
2031
2032        let pending_completion_id = post_inc(&mut self.completion_count);
2033
2034        let task = cx.spawn({
2035            async move |this, cx| {
2036                let stream = model.stream_completion(request, &cx);
2037                let assistant_message_id = assistant_message.id;
2038                let mut response_latency = None;
2039                let stream_completion = async {
2040                    let request_start = Instant::now();
2041                    let mut events = stream.await?;
2042                    let mut stop_reason = StopReason::EndTurn;
2043                    let mut thought_process_stack = Vec::new();
2044
2045                    const THOUGHT_PROCESS_START_MARKER: &str = "<think>\n";
2046                    const THOUGHT_PROCESS_END_MARKER: &str = "\n</think>";
2047
2048                    while let Some(event) = events.next().await {
2049                        if response_latency.is_none() {
2050                            response_latency = Some(request_start.elapsed());
2051                        }
2052                        let event = event?;
2053
2054                        let mut context_event = None;
2055                        let mut thought_process_output_section = None;
2056
2057                        this.update(cx, |this, cx| {
2058                            let message_ix = this
2059                                .message_anchors
2060                                .iter()
2061                                .position(|message| message.id == assistant_message_id)?;
2062                            this.buffer.update(cx, |buffer, cx| {
2063                                let message_old_end_offset = this.message_anchors[message_ix + 1..]
2064                                    .iter()
2065                                    .find(|message| message.start.is_valid(buffer))
2066                                    .map_or(buffer.len(), |message| {
2067                                        message.start.to_offset(buffer).saturating_sub(1)
2068                                    });
2069
2070                                match event {
2071                                    LanguageModelCompletionEvent::StatusUpdate { .. } => {}
2072                                    LanguageModelCompletionEvent::StartMessage { .. } => {}
2073                                    LanguageModelCompletionEvent::Stop(reason) => {
2074                                        stop_reason = reason;
2075                                    }
2076                                    LanguageModelCompletionEvent::Thinking { text: chunk, .. } => {
2077                                        if thought_process_stack.is_empty() {
2078                                            let start =
2079                                                buffer.anchor_before(message_old_end_offset);
2080                                            thought_process_stack.push(start);
2081                                            let chunk =
2082                                                format!("{THOUGHT_PROCESS_START_MARKER}{chunk}{THOUGHT_PROCESS_END_MARKER}");
2083                                            let chunk_len = chunk.len();
2084                                            buffer.edit(
2085                                                [(
2086                                                    message_old_end_offset..message_old_end_offset,
2087                                                    chunk,
2088                                                )],
2089                                                None,
2090                                                cx,
2091                                            );
2092                                            let end = buffer
2093                                                .anchor_before(message_old_end_offset + chunk_len);
2094                                            context_event = Some(
2095                                                ContextEvent::StartedThoughtProcess(start..end),
2096                                            );
2097                                        } else {
2098                                            // This ensures that all the thinking chunks are inserted inside the thinking tag
2099                                            let insertion_position =
2100                                                message_old_end_offset - THOUGHT_PROCESS_END_MARKER.len();
2101                                            buffer.edit(
2102                                                [(insertion_position..insertion_position, chunk)],
2103                                                None,
2104                                                cx,
2105                                            );
2106                                        }
2107                                    }
2108                                    LanguageModelCompletionEvent::Text(mut chunk) => {
2109                                        if let Some(start) = thought_process_stack.pop() {
2110                                            let end = buffer.anchor_before(message_old_end_offset);
2111                                            context_event =
2112                                                Some(ContextEvent::EndedThoughtProcess(end));
2113                                            thought_process_output_section =
2114                                                Some(ThoughtProcessOutputSection {
2115                                                    range: start..end,
2116                                                });
2117                                            chunk.insert_str(0, "\n\n");
2118                                        }
2119
2120                                        buffer.edit(
2121                                            [(
2122                                                message_old_end_offset..message_old_end_offset,
2123                                                chunk,
2124                                            )],
2125                                            None,
2126                                            cx,
2127                                        );
2128                                    }
2129                                    LanguageModelCompletionEvent::ToolUse(_) |
2130                                    LanguageModelCompletionEvent::UsageUpdate(_)  => {}
2131                                }
2132                            });
2133
2134                            if let Some(section) = thought_process_output_section.take() {
2135                                this.insert_thought_process_output_section(section, cx);
2136                            }
2137                            if let Some(context_event) = context_event.take() {
2138                                cx.emit(context_event);
2139                            }
2140
2141                            cx.emit(ContextEvent::StreamedCompletion);
2142
2143                            Some(())
2144                        })?;
2145                        smol::future::yield_now().await;
2146                    }
2147                    this.update(cx, |this, cx| {
2148                        this.pending_completions
2149                            .retain(|completion| completion.id != pending_completion_id);
2150                        this.summarize(false, cx);
2151                        this.update_cache_status_for_completion(cx);
2152                    })?;
2153
2154                    anyhow::Ok(stop_reason)
2155                };
2156
2157                let result = stream_completion.await;
2158
2159                this.update(cx, |this, cx| {
2160                    let error_message = if let Some(error) = result.as_ref().err() {
2161                        if error.is::<PaymentRequiredError>() {
2162                            cx.emit(ContextEvent::ShowPaymentRequiredError);
2163                            this.update_metadata(assistant_message_id, cx, |metadata| {
2164                                metadata.status = MessageStatus::Canceled;
2165                            });
2166                            Some(error.to_string())
2167                        } else {
2168                            let error_message = error
2169                                .chain()
2170                                .map(|err| err.to_string())
2171                                .collect::<Vec<_>>()
2172                                .join("\n");
2173                            cx.emit(ContextEvent::ShowAssistError(SharedString::from(
2174                                error_message.clone(),
2175                            )));
2176                            this.update_metadata(assistant_message_id, cx, |metadata| {
2177                                metadata.status =
2178                                    MessageStatus::Error(SharedString::from(error_message.clone()));
2179                            });
2180                            Some(error_message)
2181                        }
2182                    } else {
2183                        this.update_metadata(assistant_message_id, cx, |metadata| {
2184                            metadata.status = MessageStatus::Done;
2185                        });
2186                        None
2187                    };
2188
2189                    let language_name = this
2190                        .buffer
2191                        .read(cx)
2192                        .language()
2193                        .map(|language| language.name());
2194                    report_assistant_event(
2195                        AssistantEventData {
2196                            conversation_id: Some(this.id.0.clone()),
2197                            kind: AssistantKind::Panel,
2198                            phase: AssistantPhase::Response,
2199                            message_id: None,
2200                            model: model.telemetry_id(),
2201                            model_provider: model.provider_id().to_string(),
2202                            response_latency,
2203                            error_message,
2204                            language_name: language_name.map(|name| name.to_proto()),
2205                        },
2206                        this.telemetry.clone(),
2207                        cx.http_client(),
2208                        model.api_key(cx),
2209                        cx.background_executor(),
2210                    );
2211
2212                    if let Ok(stop_reason) = result {
2213                        match stop_reason {
2214                            StopReason::ToolUse => {}
2215                            StopReason::EndTurn => {}
2216                            StopReason::MaxTokens => {}
2217                            StopReason::Refusal => {}
2218                        }
2219                    }
2220                })
2221                .ok();
2222            }
2223        });
2224
2225        self.pending_completions.push(PendingCompletion {
2226            id: pending_completion_id,
2227            assistant_message_id: assistant_message.id,
2228            _task: task,
2229        });
2230
2231        Some(user_message)
2232    }
2233
2234    pub fn to_xml(&self, cx: &App) -> String {
2235        let mut output = String::new();
2236        let buffer = self.buffer.read(cx);
2237        for message in self.messages(cx) {
2238            if message.status != MessageStatus::Done {
2239                continue;
2240            }
2241
2242            writeln!(&mut output, "<{}>", message.role).unwrap();
2243            for chunk in buffer.text_for_range(message.offset_range) {
2244                output.push_str(chunk);
2245            }
2246            if !output.ends_with('\n') {
2247                output.push('\n');
2248            }
2249            writeln!(&mut output, "</{}>", message.role).unwrap();
2250        }
2251        output
2252    }
2253
2254    pub fn to_completion_request(
2255        &self,
2256        model: Option<&Arc<dyn LanguageModel>>,
2257        cx: &App,
2258    ) -> LanguageModelRequest {
2259        let buffer = self.buffer.read(cx);
2260
2261        let mut contents = self.contents(cx).peekable();
2262
2263        fn collect_text_content(buffer: &Buffer, range: Range<usize>) -> Option<String> {
2264            let text: String = buffer.text_for_range(range.clone()).collect();
2265            if text.trim().is_empty() {
2266                None
2267            } else {
2268                Some(text)
2269            }
2270        }
2271
2272        let mut completion_request = LanguageModelRequest {
2273            thread_id: None,
2274            prompt_id: None,
2275            mode: None,
2276            messages: Vec::new(),
2277            tools: Vec::new(),
2278            tool_choice: None,
2279            stop: Vec::new(),
2280            temperature: model.and_then(|model| AgentSettings::temperature_for_model(model, cx)),
2281        };
2282        for message in self.messages(cx) {
2283            if message.status != MessageStatus::Done {
2284                continue;
2285            }
2286
2287            let mut offset = message.offset_range.start;
2288            let mut request_message = LanguageModelRequestMessage {
2289                role: message.role,
2290                content: Vec::new(),
2291                cache: message
2292                    .cache
2293                    .as_ref()
2294                    .map_or(false, |cache| cache.is_anchor),
2295            };
2296
2297            while let Some(content) = contents.peek() {
2298                if content
2299                    .range()
2300                    .end
2301                    .cmp(&message.anchor_range.end, buffer)
2302                    .is_lt()
2303                {
2304                    let content = contents.next().unwrap();
2305                    let range = content.range().to_offset(buffer);
2306                    request_message.content.extend(
2307                        collect_text_content(buffer, offset..range.start).map(MessageContent::Text),
2308                    );
2309
2310                    match content {
2311                        Content::Image { image, .. } => {
2312                            if let Some(image) = image.clone().now_or_never().flatten() {
2313                                request_message
2314                                    .content
2315                                    .push(language_model::MessageContent::Image(image));
2316                            }
2317                        }
2318                    }
2319
2320                    offset = range.end;
2321                } else {
2322                    break;
2323                }
2324            }
2325
2326            request_message.content.extend(
2327                collect_text_content(buffer, offset..message.offset_range.end)
2328                    .map(MessageContent::Text),
2329            );
2330
2331            if !request_message.contents_empty() {
2332                completion_request.messages.push(request_message);
2333            }
2334        }
2335        let supports_max_mode = if let Some(model) = model {
2336            model.supports_max_mode()
2337        } else {
2338            false
2339        };
2340
2341        if supports_max_mode {
2342            completion_request.mode = Some(self.completion_mode.into());
2343        }
2344        completion_request
2345    }
2346
2347    pub fn cancel_last_assist(&mut self, cx: &mut Context<Self>) -> bool {
2348        if let Some(pending_completion) = self.pending_completions.pop() {
2349            self.update_metadata(pending_completion.assistant_message_id, cx, |metadata| {
2350                if metadata.status == MessageStatus::Pending {
2351                    metadata.status = MessageStatus::Canceled;
2352                }
2353            });
2354            true
2355        } else {
2356            false
2357        }
2358    }
2359
2360    pub fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut Context<Self>) {
2361        for id in &ids {
2362            if let Some(metadata) = self.messages_metadata.get(id) {
2363                let role = metadata.role.cycle();
2364                self.update_metadata(*id, cx, |metadata| metadata.role = role);
2365            }
2366        }
2367
2368        self.message_roles_updated(ids, cx);
2369    }
2370
2371    fn message_roles_updated(&mut self, ids: HashSet<MessageId>, cx: &mut Context<Self>) {
2372        let mut ranges = Vec::new();
2373        for message in self.messages(cx) {
2374            if ids.contains(&message.id) {
2375                ranges.push(message.anchor_range.clone());
2376            }
2377        }
2378    }
2379
2380    pub fn update_metadata(
2381        &mut self,
2382        id: MessageId,
2383        cx: &mut Context<Self>,
2384        f: impl FnOnce(&mut MessageMetadata),
2385    ) {
2386        let version = self.version.clone();
2387        let timestamp = self.next_timestamp();
2388        if let Some(metadata) = self.messages_metadata.get_mut(&id) {
2389            f(metadata);
2390            metadata.timestamp = timestamp;
2391            let operation = ContextOperation::UpdateMessage {
2392                message_id: id,
2393                metadata: metadata.clone(),
2394                version,
2395            };
2396            self.push_op(operation, cx);
2397            cx.emit(ContextEvent::MessagesEdited);
2398            cx.notify();
2399        }
2400    }
2401
2402    pub fn insert_message_after(
2403        &mut self,
2404        message_id: MessageId,
2405        role: Role,
2406        status: MessageStatus,
2407        cx: &mut Context<Self>,
2408    ) -> Option<MessageAnchor> {
2409        if let Some(prev_message_ix) = self
2410            .message_anchors
2411            .iter()
2412            .position(|message| message.id == message_id)
2413        {
2414            // Find the next valid message after the one we were given.
2415            let mut next_message_ix = prev_message_ix + 1;
2416            while let Some(next_message) = self.message_anchors.get(next_message_ix) {
2417                if next_message.start.is_valid(self.buffer.read(cx)) {
2418                    break;
2419                }
2420                next_message_ix += 1;
2421            }
2422
2423            let buffer = self.buffer.read(cx);
2424            let offset = self
2425                .message_anchors
2426                .get(next_message_ix)
2427                .map_or(buffer.len(), |message| {
2428                    buffer.clip_offset(message.start.to_offset(buffer) - 1, Bias::Left)
2429                });
2430            Some(self.insert_message_at_offset(offset, role, status, cx))
2431        } else {
2432            None
2433        }
2434    }
2435
2436    fn insert_message_at_offset(
2437        &mut self,
2438        offset: usize,
2439        role: Role,
2440        status: MessageStatus,
2441        cx: &mut Context<Self>,
2442    ) -> MessageAnchor {
2443        let start = self.buffer.update(cx, |buffer, cx| {
2444            buffer.edit([(offset..offset, "\n")], None, cx);
2445            buffer.anchor_before(offset + 1)
2446        });
2447
2448        let version = self.version.clone();
2449        let anchor = MessageAnchor {
2450            id: MessageId(self.next_timestamp()),
2451            start,
2452        };
2453        let metadata = MessageMetadata {
2454            role,
2455            status,
2456            timestamp: anchor.id.0,
2457            cache: None,
2458        };
2459        self.insert_message(anchor.clone(), metadata.clone(), cx);
2460        self.push_op(
2461            ContextOperation::InsertMessage {
2462                anchor: anchor.clone(),
2463                metadata,
2464                version,
2465            },
2466            cx,
2467        );
2468        anchor
2469    }
2470
2471    pub fn insert_content(&mut self, content: Content, cx: &mut Context<Self>) {
2472        let buffer = self.buffer.read(cx);
2473        let insertion_ix = match self
2474            .contents
2475            .binary_search_by(|probe| probe.cmp(&content, buffer))
2476        {
2477            Ok(ix) => {
2478                self.contents.remove(ix);
2479                ix
2480            }
2481            Err(ix) => ix,
2482        };
2483        self.contents.insert(insertion_ix, content);
2484        cx.emit(ContextEvent::MessagesEdited);
2485    }
2486
2487    pub fn contents<'a>(&'a self, cx: &'a App) -> impl 'a + Iterator<Item = Content> {
2488        let buffer = self.buffer.read(cx);
2489        self.contents
2490            .iter()
2491            .filter(|content| {
2492                let range = content.range();
2493                range.start.is_valid(buffer) && range.end.is_valid(buffer)
2494            })
2495            .cloned()
2496    }
2497
2498    pub fn split_message(
2499        &mut self,
2500        range: Range<usize>,
2501        cx: &mut Context<Self>,
2502    ) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
2503        let start_message = self.message_for_offset(range.start, cx);
2504        let end_message = self.message_for_offset(range.end, cx);
2505        if let Some((start_message, end_message)) = start_message.zip(end_message) {
2506            // Prevent splitting when range spans multiple messages.
2507            if start_message.id != end_message.id {
2508                return (None, None);
2509            }
2510
2511            let message = start_message;
2512            let role = message.role;
2513            let mut edited_buffer = false;
2514
2515            let mut suffix_start = None;
2516
2517            // TODO: why did this start panicking?
2518            if range.start > message.offset_range.start
2519                && range.end < message.offset_range.end.saturating_sub(1)
2520            {
2521                if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
2522                    suffix_start = Some(range.end + 1);
2523                } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
2524                    suffix_start = Some(range.end);
2525                }
2526            }
2527
2528            let version = self.version.clone();
2529            let suffix = if let Some(suffix_start) = suffix_start {
2530                MessageAnchor {
2531                    id: MessageId(self.next_timestamp()),
2532                    start: self.buffer.read(cx).anchor_before(suffix_start),
2533                }
2534            } else {
2535                self.buffer.update(cx, |buffer, cx| {
2536                    buffer.edit([(range.end..range.end, "\n")], None, cx);
2537                });
2538                edited_buffer = true;
2539                MessageAnchor {
2540                    id: MessageId(self.next_timestamp()),
2541                    start: self.buffer.read(cx).anchor_before(range.end + 1),
2542                }
2543            };
2544
2545            let suffix_metadata = MessageMetadata {
2546                role,
2547                status: MessageStatus::Done,
2548                timestamp: suffix.id.0,
2549                cache: None,
2550            };
2551            self.insert_message(suffix.clone(), suffix_metadata.clone(), cx);
2552            self.push_op(
2553                ContextOperation::InsertMessage {
2554                    anchor: suffix.clone(),
2555                    metadata: suffix_metadata,
2556                    version,
2557                },
2558                cx,
2559            );
2560
2561            let new_messages =
2562                if range.start == range.end || range.start == message.offset_range.start {
2563                    (None, Some(suffix))
2564                } else {
2565                    let mut prefix_end = None;
2566                    if range.start > message.offset_range.start
2567                        && range.end < message.offset_range.end - 1
2568                    {
2569                        if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
2570                            prefix_end = Some(range.start + 1);
2571                        } else if self.buffer.read(cx).reversed_chars_at(range.start).next()
2572                            == Some('\n')
2573                        {
2574                            prefix_end = Some(range.start);
2575                        }
2576                    }
2577
2578                    let version = self.version.clone();
2579                    let selection = if let Some(prefix_end) = prefix_end {
2580                        MessageAnchor {
2581                            id: MessageId(self.next_timestamp()),
2582                            start: self.buffer.read(cx).anchor_before(prefix_end),
2583                        }
2584                    } else {
2585                        self.buffer.update(cx, |buffer, cx| {
2586                            buffer.edit([(range.start..range.start, "\n")], None, cx)
2587                        });
2588                        edited_buffer = true;
2589                        MessageAnchor {
2590                            id: MessageId(self.next_timestamp()),
2591                            start: self.buffer.read(cx).anchor_before(range.end + 1),
2592                        }
2593                    };
2594
2595                    let selection_metadata = MessageMetadata {
2596                        role,
2597                        status: MessageStatus::Done,
2598                        timestamp: selection.id.0,
2599                        cache: None,
2600                    };
2601                    self.insert_message(selection.clone(), selection_metadata.clone(), cx);
2602                    self.push_op(
2603                        ContextOperation::InsertMessage {
2604                            anchor: selection.clone(),
2605                            metadata: selection_metadata,
2606                            version,
2607                        },
2608                        cx,
2609                    );
2610
2611                    (Some(selection), Some(suffix))
2612                };
2613
2614            if !edited_buffer {
2615                cx.emit(ContextEvent::MessagesEdited);
2616            }
2617            new_messages
2618        } else {
2619            (None, None)
2620        }
2621    }
2622
2623    fn insert_message(
2624        &mut self,
2625        new_anchor: MessageAnchor,
2626        new_metadata: MessageMetadata,
2627        cx: &mut Context<Self>,
2628    ) {
2629        cx.emit(ContextEvent::MessagesEdited);
2630
2631        self.messages_metadata.insert(new_anchor.id, new_metadata);
2632
2633        let buffer = self.buffer.read(cx);
2634        let insertion_ix = self
2635            .message_anchors
2636            .iter()
2637            .position(|anchor| {
2638                let comparison = new_anchor.start.cmp(&anchor.start, buffer);
2639                comparison.is_lt() || (comparison.is_eq() && new_anchor.id > anchor.id)
2640            })
2641            .unwrap_or(self.message_anchors.len());
2642        self.message_anchors.insert(insertion_ix, new_anchor);
2643    }
2644
2645    pub fn summarize(&mut self, mut replace_old: bool, cx: &mut Context<Self>) {
2646        let Some(model) = LanguageModelRegistry::read_global(cx).default_model() else {
2647            return;
2648        };
2649
2650        if replace_old || (self.message_anchors.len() >= 2 && self.summary.is_pending()) {
2651            if !model.provider.is_authenticated(cx) {
2652                return;
2653            }
2654
2655            let mut request = self.to_completion_request(Some(&model.model), cx);
2656            request.messages.push(LanguageModelRequestMessage {
2657                role: Role::User,
2658                content: vec![
2659                    "Generate a concise 3-7 word title for this conversation, omitting punctuation. Go straight to the title, without any preamble and prefix like `Here's a concise suggestion:...` or `Title:`"
2660                        .into(),
2661                ],
2662                cache: false,
2663            });
2664
2665            // If there is no summary, it is set with `done: false` so that "Loading Summary…" can
2666            // be displayed.
2667            match self.summary {
2668                ContextSummary::Pending | ContextSummary::Error => {
2669                    self.summary = ContextSummary::Content(ContextSummaryContent {
2670                        text: "".to_string(),
2671                        done: false,
2672                        timestamp: clock::Lamport::default(),
2673                    });
2674                    replace_old = true;
2675                }
2676                ContextSummary::Content(_) => {}
2677            }
2678
2679            self.summary_task = cx.spawn(async move |this, cx| {
2680                let result = async {
2681                    let stream = model.model.stream_completion_text(request, &cx);
2682                    let mut messages = stream.await?;
2683
2684                    let mut replaced = !replace_old;
2685                    while let Some(message) = messages.stream.next().await {
2686                        let text = message?;
2687                        let mut lines = text.lines();
2688                        this.update(cx, |this, cx| {
2689                            let version = this.version.clone();
2690                            let timestamp = this.next_timestamp();
2691                            let summary = this.summary.content_or_set_empty();
2692                            if !replaced && replace_old {
2693                                summary.text.clear();
2694                                replaced = true;
2695                            }
2696                            summary.text.extend(lines.next());
2697                            summary.timestamp = timestamp;
2698                            let operation = ContextOperation::UpdateSummary {
2699                                summary: summary.clone(),
2700                                version,
2701                            };
2702                            this.push_op(operation, cx);
2703                            cx.emit(ContextEvent::SummaryChanged);
2704                            cx.emit(ContextEvent::SummaryGenerated);
2705                        })?;
2706
2707                        // Stop if the LLM generated multiple lines.
2708                        if lines.next().is_some() {
2709                            break;
2710                        }
2711                    }
2712
2713                    this.read_with(cx, |this, _cx| {
2714                        if let Some(summary) = this.summary.content() {
2715                            if summary.text.is_empty() {
2716                                bail!("Model generated an empty summary");
2717                            }
2718                        }
2719                        Ok(())
2720                    })??;
2721
2722                    this.update(cx, |this, cx| {
2723                        let version = this.version.clone();
2724                        let timestamp = this.next_timestamp();
2725                        if let Some(summary) = this.summary.content_as_mut() {
2726                            summary.done = true;
2727                            summary.timestamp = timestamp;
2728                            let operation = ContextOperation::UpdateSummary {
2729                                summary: summary.clone(),
2730                                version,
2731                            };
2732                            this.push_op(operation, cx);
2733                            cx.emit(ContextEvent::SummaryChanged);
2734                            cx.emit(ContextEvent::SummaryGenerated);
2735                        }
2736                    })?;
2737
2738                    anyhow::Ok(())
2739                }
2740                .await;
2741
2742                if let Err(err) = result {
2743                    this.update(cx, |this, cx| {
2744                        this.summary = ContextSummary::Error;
2745                        cx.emit(ContextEvent::SummaryChanged);
2746                    })
2747                    .log_err();
2748                    log::error!("Error generating context summary: {}", err);
2749                }
2750
2751                Some(())
2752            });
2753        }
2754    }
2755
2756    fn message_for_offset(&self, offset: usize, cx: &App) -> Option<Message> {
2757        self.messages_for_offsets([offset], cx).pop()
2758    }
2759
2760    pub fn messages_for_offsets(
2761        &self,
2762        offsets: impl IntoIterator<Item = usize>,
2763        cx: &App,
2764    ) -> Vec<Message> {
2765        let mut result = Vec::new();
2766
2767        let mut messages = self.messages(cx).peekable();
2768        let mut offsets = offsets.into_iter().peekable();
2769        let mut current_message = messages.next();
2770        while let Some(offset) = offsets.next() {
2771            // Locate the message that contains the offset.
2772            while current_message.as_ref().map_or(false, |message| {
2773                !message.offset_range.contains(&offset) && messages.peek().is_some()
2774            }) {
2775                current_message = messages.next();
2776            }
2777            let Some(message) = current_message.as_ref() else {
2778                break;
2779            };
2780
2781            // Skip offsets that are in the same message.
2782            while offsets.peek().map_or(false, |offset| {
2783                message.offset_range.contains(offset) || messages.peek().is_none()
2784            }) {
2785                offsets.next();
2786            }
2787
2788            result.push(message.clone());
2789        }
2790        result
2791    }
2792
2793    fn messages_from_anchors<'a>(
2794        &'a self,
2795        message_anchors: impl Iterator<Item = &'a MessageAnchor> + 'a,
2796        cx: &'a App,
2797    ) -> impl 'a + Iterator<Item = Message> {
2798        let buffer = self.buffer.read(cx);
2799
2800        Self::messages_from_iters(buffer, &self.messages_metadata, message_anchors.enumerate())
2801    }
2802
2803    pub fn messages<'a>(&'a self, cx: &'a App) -> impl 'a + Iterator<Item = Message> {
2804        self.messages_from_anchors(self.message_anchors.iter(), cx)
2805    }
2806
2807    pub fn messages_from_iters<'a>(
2808        buffer: &'a Buffer,
2809        metadata: &'a HashMap<MessageId, MessageMetadata>,
2810        messages: impl Iterator<Item = (usize, &'a MessageAnchor)> + 'a,
2811    ) -> impl 'a + Iterator<Item = Message> {
2812        let mut messages = messages.peekable();
2813
2814        iter::from_fn(move || {
2815            if let Some((start_ix, message_anchor)) = messages.next() {
2816                let metadata = metadata.get(&message_anchor.id)?;
2817
2818                let message_start = message_anchor.start.to_offset(buffer);
2819                let mut message_end = None;
2820                let mut end_ix = start_ix;
2821                while let Some((_, next_message)) = messages.peek() {
2822                    if next_message.start.is_valid(buffer) {
2823                        message_end = Some(next_message.start);
2824                        break;
2825                    } else {
2826                        end_ix += 1;
2827                        messages.next();
2828                    }
2829                }
2830                let message_end_anchor = message_end.unwrap_or(language::Anchor::MAX);
2831                let message_end = message_end_anchor.to_offset(buffer);
2832
2833                return Some(Message {
2834                    index_range: start_ix..end_ix,
2835                    offset_range: message_start..message_end,
2836                    anchor_range: message_anchor.start..message_end_anchor,
2837                    id: message_anchor.id,
2838                    role: metadata.role,
2839                    status: metadata.status.clone(),
2840                    cache: metadata.cache.clone(),
2841                });
2842            }
2843            None
2844        })
2845    }
2846
2847    pub fn save(
2848        &mut self,
2849        debounce: Option<Duration>,
2850        fs: Arc<dyn Fs>,
2851        cx: &mut Context<AssistantContext>,
2852    ) {
2853        if self.replica_id() != ReplicaId::default() {
2854            // Prevent saving a remote context for now.
2855            return;
2856        }
2857
2858        self.pending_save = cx.spawn(async move |this, cx| {
2859            if let Some(debounce) = debounce {
2860                cx.background_executor().timer(debounce).await;
2861            }
2862
2863            let (old_path, summary) = this.read_with(cx, |this, _| {
2864                let path = this.path.clone();
2865                let summary = if let Some(summary) = this.summary.content() {
2866                    if summary.done {
2867                        Some(summary.text.clone())
2868                    } else {
2869                        None
2870                    }
2871                } else {
2872                    None
2873                };
2874                (path, summary)
2875            })?;
2876
2877            if let Some(summary) = summary {
2878                let context = this.read_with(cx, |this, cx| this.serialize(cx))?;
2879                let mut discriminant = 1;
2880                let mut new_path;
2881                loop {
2882                    new_path = contexts_dir().join(&format!(
2883                        "{} - {}.zed.json",
2884                        summary.trim(),
2885                        discriminant
2886                    ));
2887                    if fs.is_file(&new_path).await {
2888                        discriminant += 1;
2889                    } else {
2890                        break;
2891                    }
2892                }
2893
2894                fs.create_dir(contexts_dir().as_ref()).await?;
2895                fs.atomic_write(new_path.clone(), serde_json::to_string(&context).unwrap())
2896                    .await?;
2897                if let Some(old_path) = old_path {
2898                    if new_path.as_path() != old_path.as_ref() {
2899                        fs.remove_file(
2900                            &old_path,
2901                            RemoveOptions {
2902                                recursive: false,
2903                                ignore_if_not_exists: true,
2904                            },
2905                        )
2906                        .await?;
2907                    }
2908                }
2909
2910                this.update(cx, |this, _| this.path = Some(new_path.into()))?;
2911            }
2912
2913            Ok(())
2914        });
2915    }
2916
2917    pub fn set_custom_summary(&mut self, custom_summary: String, cx: &mut Context<Self>) {
2918        let timestamp = self.next_timestamp();
2919        let summary = self.summary.content_or_set_empty();
2920        summary.timestamp = timestamp;
2921        summary.done = true;
2922        summary.text = custom_summary;
2923        cx.emit(ContextEvent::SummaryChanged);
2924    }
2925}
2926
2927#[derive(Debug, Default)]
2928pub struct ContextVersion {
2929    context: clock::Global,
2930    buffer: clock::Global,
2931}
2932
2933impl ContextVersion {
2934    pub fn from_proto(proto: &proto::ContextVersion) -> Self {
2935        Self {
2936            context: language::proto::deserialize_version(&proto.context_version),
2937            buffer: language::proto::deserialize_version(&proto.buffer_version),
2938        }
2939    }
2940
2941    pub fn to_proto(&self, context_id: ContextId) -> proto::ContextVersion {
2942        proto::ContextVersion {
2943            context_id: context_id.to_proto(),
2944            context_version: language::proto::serialize_version(&self.context),
2945            buffer_version: language::proto::serialize_version(&self.buffer),
2946        }
2947    }
2948}
2949
2950#[derive(Debug, Clone)]
2951pub struct ParsedSlashCommand {
2952    pub name: String,
2953    pub arguments: SmallVec<[String; 3]>,
2954    pub status: PendingSlashCommandStatus,
2955    pub source_range: Range<language::Anchor>,
2956}
2957
2958#[derive(Debug)]
2959pub struct InvokedSlashCommand {
2960    pub name: SharedString,
2961    pub range: Range<language::Anchor>,
2962    pub run_commands_in_ranges: Vec<Range<language::Anchor>>,
2963    pub status: InvokedSlashCommandStatus,
2964    pub transaction: Option<language::TransactionId>,
2965    timestamp: clock::Lamport,
2966}
2967
2968#[derive(Debug)]
2969pub enum InvokedSlashCommandStatus {
2970    Running(Task<()>),
2971    Error(SharedString),
2972    Finished,
2973}
2974
2975#[derive(Debug, Clone)]
2976pub enum PendingSlashCommandStatus {
2977    Idle,
2978    Running { _task: Shared<Task<()>> },
2979    Error(String),
2980}
2981
2982#[derive(Debug, Clone)]
2983pub struct PendingToolUse {
2984    pub id: LanguageModelToolUseId,
2985    pub name: String,
2986    pub input: serde_json::Value,
2987    pub status: PendingToolUseStatus,
2988    pub source_range: Range<language::Anchor>,
2989}
2990
2991#[derive(Debug, Clone)]
2992pub enum PendingToolUseStatus {
2993    Idle,
2994    Running { _task: Shared<Task<()>> },
2995    Error(String),
2996}
2997
2998impl PendingToolUseStatus {
2999    pub fn is_idle(&self) -> bool {
3000        matches!(self, PendingToolUseStatus::Idle)
3001    }
3002}
3003
3004#[derive(Serialize, Deserialize)]
3005pub struct SavedMessage {
3006    pub id: MessageId,
3007    pub start: usize,
3008    pub metadata: MessageMetadata,
3009}
3010
3011#[derive(Serialize, Deserialize)]
3012pub struct SavedContext {
3013    pub id: Option<ContextId>,
3014    pub zed: String,
3015    pub version: String,
3016    pub text: String,
3017    pub messages: Vec<SavedMessage>,
3018    pub summary: String,
3019    pub slash_command_output_sections:
3020        Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
3021    #[serde(default)]
3022    pub thought_process_output_sections: Vec<ThoughtProcessOutputSection<usize>>,
3023}
3024
3025impl SavedContext {
3026    pub const VERSION: &'static str = "0.4.0";
3027
3028    pub fn from_json(json: &str) -> Result<Self> {
3029        let saved_context_json = serde_json::from_str::<serde_json::Value>(json)?;
3030        match saved_context_json
3031            .get("version")
3032            .context("version not found")?
3033        {
3034            serde_json::Value::String(version) => match version.as_str() {
3035                SavedContext::VERSION => {
3036                    Ok(serde_json::from_value::<SavedContext>(saved_context_json)?)
3037                }
3038                SavedContextV0_3_0::VERSION => {
3039                    let saved_context =
3040                        serde_json::from_value::<SavedContextV0_3_0>(saved_context_json)?;
3041                    Ok(saved_context.upgrade())
3042                }
3043                SavedContextV0_2_0::VERSION => {
3044                    let saved_context =
3045                        serde_json::from_value::<SavedContextV0_2_0>(saved_context_json)?;
3046                    Ok(saved_context.upgrade())
3047                }
3048                SavedContextV0_1_0::VERSION => {
3049                    let saved_context =
3050                        serde_json::from_value::<SavedContextV0_1_0>(saved_context_json)?;
3051                    Ok(saved_context.upgrade())
3052                }
3053                _ => anyhow::bail!("unrecognized saved context version: {version:?}"),
3054            },
3055            _ => anyhow::bail!("version not found on saved context"),
3056        }
3057    }
3058
3059    fn into_ops(
3060        self,
3061        buffer: &Entity<Buffer>,
3062        cx: &mut Context<AssistantContext>,
3063    ) -> Vec<ContextOperation> {
3064        let mut operations = Vec::new();
3065        let mut version = clock::Global::new();
3066        let mut next_timestamp = clock::Lamport::new(ReplicaId::default());
3067
3068        let mut first_message_metadata = None;
3069        for message in self.messages {
3070            if message.id == MessageId(clock::Lamport::default()) {
3071                first_message_metadata = Some(message.metadata);
3072            } else {
3073                operations.push(ContextOperation::InsertMessage {
3074                    anchor: MessageAnchor {
3075                        id: message.id,
3076                        start: buffer.read(cx).anchor_before(message.start),
3077                    },
3078                    metadata: MessageMetadata {
3079                        role: message.metadata.role,
3080                        status: message.metadata.status,
3081                        timestamp: message.metadata.timestamp,
3082                        cache: None,
3083                    },
3084                    version: version.clone(),
3085                });
3086                version.observe(message.id.0);
3087                next_timestamp.observe(message.id.0);
3088            }
3089        }
3090
3091        if let Some(metadata) = first_message_metadata {
3092            let timestamp = next_timestamp.tick();
3093            operations.push(ContextOperation::UpdateMessage {
3094                message_id: MessageId(clock::Lamport::default()),
3095                metadata: MessageMetadata {
3096                    role: metadata.role,
3097                    status: metadata.status,
3098                    timestamp,
3099                    cache: None,
3100                },
3101                version: version.clone(),
3102            });
3103            version.observe(timestamp);
3104        }
3105
3106        let buffer = buffer.read(cx);
3107        for section in self.slash_command_output_sections {
3108            let timestamp = next_timestamp.tick();
3109            operations.push(ContextOperation::SlashCommandOutputSectionAdded {
3110                timestamp,
3111                section: SlashCommandOutputSection {
3112                    range: buffer.anchor_after(section.range.start)
3113                        ..buffer.anchor_before(section.range.end),
3114                    icon: section.icon,
3115                    label: section.label,
3116                    metadata: section.metadata,
3117                },
3118                version: version.clone(),
3119            });
3120
3121            version.observe(timestamp);
3122        }
3123
3124        for section in self.thought_process_output_sections {
3125            let timestamp = next_timestamp.tick();
3126            operations.push(ContextOperation::ThoughtProcessOutputSectionAdded {
3127                timestamp,
3128                section: ThoughtProcessOutputSection {
3129                    range: buffer.anchor_after(section.range.start)
3130                        ..buffer.anchor_before(section.range.end),
3131                },
3132                version: version.clone(),
3133            });
3134
3135            version.observe(timestamp);
3136        }
3137
3138        let timestamp = next_timestamp.tick();
3139        operations.push(ContextOperation::UpdateSummary {
3140            summary: ContextSummaryContent {
3141                text: self.summary,
3142                done: true,
3143                timestamp,
3144            },
3145            version: version.clone(),
3146        });
3147        version.observe(timestamp);
3148
3149        operations
3150    }
3151}
3152
3153#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
3154struct SavedMessageIdPreV0_4_0(usize);
3155
3156#[derive(Serialize, Deserialize)]
3157struct SavedMessagePreV0_4_0 {
3158    id: SavedMessageIdPreV0_4_0,
3159    start: usize,
3160}
3161
3162#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
3163struct SavedMessageMetadataPreV0_4_0 {
3164    role: Role,
3165    status: MessageStatus,
3166}
3167
3168#[derive(Serialize, Deserialize)]
3169struct SavedContextV0_3_0 {
3170    id: Option<ContextId>,
3171    zed: String,
3172    version: String,
3173    text: String,
3174    messages: Vec<SavedMessagePreV0_4_0>,
3175    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3176    summary: String,
3177    slash_command_output_sections: Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
3178}
3179
3180impl SavedContextV0_3_0 {
3181    const VERSION: &'static str = "0.3.0";
3182
3183    fn upgrade(self) -> SavedContext {
3184        SavedContext {
3185            id: self.id,
3186            zed: self.zed,
3187            version: SavedContext::VERSION.into(),
3188            text: self.text,
3189            messages: self
3190                .messages
3191                .into_iter()
3192                .filter_map(|message| {
3193                    let metadata = self.message_metadata.get(&message.id)?;
3194                    let timestamp = clock::Lamport {
3195                        replica_id: ReplicaId::default(),
3196                        value: message.id.0 as u32,
3197                    };
3198                    Some(SavedMessage {
3199                        id: MessageId(timestamp),
3200                        start: message.start,
3201                        metadata: MessageMetadata {
3202                            role: metadata.role,
3203                            status: metadata.status.clone(),
3204                            timestamp,
3205                            cache: None,
3206                        },
3207                    })
3208                })
3209                .collect(),
3210            summary: self.summary,
3211            slash_command_output_sections: self.slash_command_output_sections,
3212            thought_process_output_sections: Vec::new(),
3213        }
3214    }
3215}
3216
3217#[derive(Serialize, Deserialize)]
3218struct SavedContextV0_2_0 {
3219    id: Option<ContextId>,
3220    zed: String,
3221    version: String,
3222    text: String,
3223    messages: Vec<SavedMessagePreV0_4_0>,
3224    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3225    summary: String,
3226}
3227
3228impl SavedContextV0_2_0 {
3229    const VERSION: &'static str = "0.2.0";
3230
3231    fn upgrade(self) -> SavedContext {
3232        SavedContextV0_3_0 {
3233            id: self.id,
3234            zed: self.zed,
3235            version: SavedContextV0_3_0::VERSION.to_string(),
3236            text: self.text,
3237            messages: self.messages,
3238            message_metadata: self.message_metadata,
3239            summary: self.summary,
3240            slash_command_output_sections: Vec::new(),
3241        }
3242        .upgrade()
3243    }
3244}
3245
3246#[derive(Serialize, Deserialize)]
3247struct SavedContextV0_1_0 {
3248    id: Option<ContextId>,
3249    zed: String,
3250    version: String,
3251    text: String,
3252    messages: Vec<SavedMessagePreV0_4_0>,
3253    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3254    summary: String,
3255    api_url: Option<String>,
3256    model: OpenAiModel,
3257}
3258
3259impl SavedContextV0_1_0 {
3260    const VERSION: &'static str = "0.1.0";
3261
3262    fn upgrade(self) -> SavedContext {
3263        SavedContextV0_2_0 {
3264            id: self.id,
3265            zed: self.zed,
3266            version: SavedContextV0_2_0::VERSION.to_string(),
3267            text: self.text,
3268            messages: self.messages,
3269            message_metadata: self.message_metadata,
3270            summary: self.summary,
3271        }
3272        .upgrade()
3273    }
3274}
3275
3276#[derive(Debug, Clone)]
3277pub struct SavedContextMetadata {
3278    pub title: String,
3279    pub path: Arc<Path>,
3280    pub mtime: chrono::DateTime<chrono::Local>,
3281}