assistant_context.rs

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