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::UsageUpdate(_)  => {}
2144                                }
2145                            });
2146
2147                            if let Some(section) = thought_process_output_section.take() {
2148                                this.insert_thought_process_output_section(section, cx);
2149                            }
2150                            if let Some(context_event) = context_event.take() {
2151                                cx.emit(context_event);
2152                            }
2153
2154                            cx.emit(ContextEvent::StreamedCompletion);
2155
2156                            Some(())
2157                        })?;
2158                        smol::future::yield_now().await;
2159                    }
2160                    this.update(cx, |this, cx| {
2161                        this.pending_completions
2162                            .retain(|completion| completion.id != pending_completion_id);
2163                        this.summarize(false, cx);
2164                        this.update_cache_status_for_completion(cx);
2165                    })?;
2166
2167                    anyhow::Ok(stop_reason)
2168                };
2169
2170                let result = stream_completion.await;
2171
2172                this.update(cx, |this, cx| {
2173                    let error_message = if let Some(error) = result.as_ref().err() {
2174                        if error.is::<PaymentRequiredError>() {
2175                            cx.emit(ContextEvent::ShowPaymentRequiredError);
2176                            this.update_metadata(assistant_message_id, cx, |metadata| {
2177                                metadata.status = MessageStatus::Canceled;
2178                            });
2179                            Some(error.to_string())
2180                        } else {
2181                            let error_message = error
2182                                .chain()
2183                                .map(|err| err.to_string())
2184                                .collect::<Vec<_>>()
2185                                .join("\n");
2186                            cx.emit(ContextEvent::ShowAssistError(SharedString::from(
2187                                error_message.clone(),
2188                            )));
2189                            this.update_metadata(assistant_message_id, cx, |metadata| {
2190                                metadata.status =
2191                                    MessageStatus::Error(SharedString::from(error_message.clone()));
2192                            });
2193                            Some(error_message)
2194                        }
2195                    } else {
2196                        this.update_metadata(assistant_message_id, cx, |metadata| {
2197                            metadata.status = MessageStatus::Done;
2198                        });
2199                        None
2200                    };
2201
2202                    let language_name = this
2203                        .buffer
2204                        .read(cx)
2205                        .language()
2206                        .map(|language| language.name());
2207                    report_assistant_event(
2208                        AssistantEventData {
2209                            conversation_id: Some(this.id.0.clone()),
2210                            kind: AssistantKind::Panel,
2211                            phase: AssistantPhase::Response,
2212                            message_id: None,
2213                            model: model.telemetry_id(),
2214                            model_provider: model.provider_id().to_string(),
2215                            response_latency,
2216                            error_message,
2217                            language_name: language_name.map(|name| name.to_proto()),
2218                        },
2219                        this.telemetry.clone(),
2220                        cx.http_client(),
2221                        model.api_key(cx),
2222                        cx.background_executor(),
2223                    );
2224
2225                    if let Ok(stop_reason) = result {
2226                        match stop_reason {
2227                            StopReason::ToolUse => {}
2228                            StopReason::EndTurn => {}
2229                            StopReason::MaxTokens => {}
2230                            StopReason::Refusal => {}
2231                        }
2232                    }
2233                })
2234                .ok();
2235            }
2236        });
2237
2238        self.pending_completions.push(PendingCompletion {
2239            id: pending_completion_id,
2240            assistant_message_id: assistant_message.id,
2241            _task: task,
2242        });
2243
2244        Some(user_message)
2245    }
2246
2247    pub fn to_xml(&self, cx: &App) -> String {
2248        let mut output = String::new();
2249        let buffer = self.buffer.read(cx);
2250        for message in self.messages(cx) {
2251            if message.status != MessageStatus::Done {
2252                continue;
2253            }
2254
2255            writeln!(&mut output, "<{}>", message.role).unwrap();
2256            for chunk in buffer.text_for_range(message.offset_range) {
2257                output.push_str(chunk);
2258            }
2259            if !output.ends_with('\n') {
2260                output.push('\n');
2261            }
2262            writeln!(&mut output, "</{}>", message.role).unwrap();
2263        }
2264        output
2265    }
2266
2267    pub fn to_completion_request(
2268        &self,
2269        model: Option<&Arc<dyn LanguageModel>>,
2270        cx: &App,
2271    ) -> LanguageModelRequest {
2272        let buffer = self.buffer.read(cx);
2273
2274        let mut contents = self.contents(cx).peekable();
2275
2276        fn collect_text_content(buffer: &Buffer, range: Range<usize>) -> Option<String> {
2277            let text: String = buffer.text_for_range(range.clone()).collect();
2278            if text.trim().is_empty() {
2279                None
2280            } else {
2281                Some(text)
2282            }
2283        }
2284
2285        let mut completion_request = LanguageModelRequest {
2286            thread_id: None,
2287            prompt_id: None,
2288            intent: Some(CompletionIntent::UserPrompt),
2289            mode: None,
2290            messages: Vec::new(),
2291            tools: Vec::new(),
2292            tool_choice: None,
2293            stop: Vec::new(),
2294            temperature: model.and_then(|model| AgentSettings::temperature_for_model(model, cx)),
2295        };
2296        for message in self.messages(cx) {
2297            if message.status != MessageStatus::Done {
2298                continue;
2299            }
2300
2301            let mut offset = message.offset_range.start;
2302            let mut request_message = LanguageModelRequestMessage {
2303                role: message.role,
2304                content: Vec::new(),
2305                cache: message
2306                    .cache
2307                    .as_ref()
2308                    .map_or(false, |cache| cache.is_anchor),
2309            };
2310
2311            while let Some(content) = contents.peek() {
2312                if content
2313                    .range()
2314                    .end
2315                    .cmp(&message.anchor_range.end, buffer)
2316                    .is_lt()
2317                {
2318                    let content = contents.next().unwrap();
2319                    let range = content.range().to_offset(buffer);
2320                    request_message.content.extend(
2321                        collect_text_content(buffer, offset..range.start).map(MessageContent::Text),
2322                    );
2323
2324                    match content {
2325                        Content::Image { image, .. } => {
2326                            if let Some(image) = image.clone().now_or_never().flatten() {
2327                                request_message
2328                                    .content
2329                                    .push(language_model::MessageContent::Image(image));
2330                            }
2331                        }
2332                    }
2333
2334                    offset = range.end;
2335                } else {
2336                    break;
2337                }
2338            }
2339
2340            request_message.content.extend(
2341                collect_text_content(buffer, offset..message.offset_range.end)
2342                    .map(MessageContent::Text),
2343            );
2344
2345            if !request_message.contents_empty() {
2346                completion_request.messages.push(request_message);
2347            }
2348        }
2349        let supports_max_mode = if let Some(model) = model {
2350            model.supports_max_mode()
2351        } else {
2352            false
2353        };
2354
2355        if supports_max_mode {
2356            completion_request.mode = Some(self.completion_mode.into());
2357        }
2358        completion_request
2359    }
2360
2361    pub fn cancel_last_assist(&mut self, cx: &mut Context<Self>) -> bool {
2362        if let Some(pending_completion) = self.pending_completions.pop() {
2363            self.update_metadata(pending_completion.assistant_message_id, cx, |metadata| {
2364                if metadata.status == MessageStatus::Pending {
2365                    metadata.status = MessageStatus::Canceled;
2366                }
2367            });
2368            true
2369        } else {
2370            false
2371        }
2372    }
2373
2374    pub fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut Context<Self>) {
2375        for id in &ids {
2376            if let Some(metadata) = self.messages_metadata.get(id) {
2377                let role = metadata.role.cycle();
2378                self.update_metadata(*id, cx, |metadata| metadata.role = role);
2379            }
2380        }
2381
2382        self.message_roles_updated(ids, cx);
2383    }
2384
2385    fn message_roles_updated(&mut self, ids: HashSet<MessageId>, cx: &mut Context<Self>) {
2386        let mut ranges = Vec::new();
2387        for message in self.messages(cx) {
2388            if ids.contains(&message.id) {
2389                ranges.push(message.anchor_range.clone());
2390            }
2391        }
2392    }
2393
2394    pub fn update_metadata(
2395        &mut self,
2396        id: MessageId,
2397        cx: &mut Context<Self>,
2398        f: impl FnOnce(&mut MessageMetadata),
2399    ) {
2400        let version = self.version.clone();
2401        let timestamp = self.next_timestamp();
2402        if let Some(metadata) = self.messages_metadata.get_mut(&id) {
2403            f(metadata);
2404            metadata.timestamp = timestamp;
2405            let operation = ContextOperation::UpdateMessage {
2406                message_id: id,
2407                metadata: metadata.clone(),
2408                version,
2409            };
2410            self.push_op(operation, cx);
2411            cx.emit(ContextEvent::MessagesEdited);
2412            cx.notify();
2413        }
2414    }
2415
2416    pub fn insert_message_after(
2417        &mut self,
2418        message_id: MessageId,
2419        role: Role,
2420        status: MessageStatus,
2421        cx: &mut Context<Self>,
2422    ) -> Option<MessageAnchor> {
2423        if let Some(prev_message_ix) = self
2424            .message_anchors
2425            .iter()
2426            .position(|message| message.id == message_id)
2427        {
2428            // Find the next valid message after the one we were given.
2429            let mut next_message_ix = prev_message_ix + 1;
2430            while let Some(next_message) = self.message_anchors.get(next_message_ix) {
2431                if next_message.start.is_valid(self.buffer.read(cx)) {
2432                    break;
2433                }
2434                next_message_ix += 1;
2435            }
2436
2437            let buffer = self.buffer.read(cx);
2438            let offset = self
2439                .message_anchors
2440                .get(next_message_ix)
2441                .map_or(buffer.len(), |message| {
2442                    buffer.clip_offset(message.start.to_offset(buffer) - 1, Bias::Left)
2443                });
2444            Some(self.insert_message_at_offset(offset, role, status, cx))
2445        } else {
2446            None
2447        }
2448    }
2449
2450    fn insert_message_at_offset(
2451        &mut self,
2452        offset: usize,
2453        role: Role,
2454        status: MessageStatus,
2455        cx: &mut Context<Self>,
2456    ) -> MessageAnchor {
2457        let start = self.buffer.update(cx, |buffer, cx| {
2458            buffer.edit([(offset..offset, "\n")], None, cx);
2459            buffer.anchor_before(offset + 1)
2460        });
2461
2462        let version = self.version.clone();
2463        let anchor = MessageAnchor {
2464            id: MessageId(self.next_timestamp()),
2465            start,
2466        };
2467        let metadata = MessageMetadata {
2468            role,
2469            status,
2470            timestamp: anchor.id.0,
2471            cache: None,
2472        };
2473        self.insert_message(anchor.clone(), metadata.clone(), cx);
2474        self.push_op(
2475            ContextOperation::InsertMessage {
2476                anchor: anchor.clone(),
2477                metadata,
2478                version,
2479            },
2480            cx,
2481        );
2482        anchor
2483    }
2484
2485    pub fn insert_content(&mut self, content: Content, cx: &mut Context<Self>) {
2486        let buffer = self.buffer.read(cx);
2487        let insertion_ix = match self
2488            .contents
2489            .binary_search_by(|probe| probe.cmp(&content, buffer))
2490        {
2491            Ok(ix) => {
2492                self.contents.remove(ix);
2493                ix
2494            }
2495            Err(ix) => ix,
2496        };
2497        self.contents.insert(insertion_ix, content);
2498        cx.emit(ContextEvent::MessagesEdited);
2499    }
2500
2501    pub fn contents<'a>(&'a self, cx: &'a App) -> impl 'a + Iterator<Item = Content> {
2502        let buffer = self.buffer.read(cx);
2503        self.contents
2504            .iter()
2505            .filter(|content| {
2506                let range = content.range();
2507                range.start.is_valid(buffer) && range.end.is_valid(buffer)
2508            })
2509            .cloned()
2510    }
2511
2512    pub fn split_message(
2513        &mut self,
2514        range: Range<usize>,
2515        cx: &mut Context<Self>,
2516    ) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
2517        let start_message = self.message_for_offset(range.start, cx);
2518        let end_message = self.message_for_offset(range.end, cx);
2519        if let Some((start_message, end_message)) = start_message.zip(end_message) {
2520            // Prevent splitting when range spans multiple messages.
2521            if start_message.id != end_message.id {
2522                return (None, None);
2523            }
2524
2525            let message = start_message;
2526            let role = message.role;
2527            let mut edited_buffer = false;
2528
2529            let mut suffix_start = None;
2530
2531            // TODO: why did this start panicking?
2532            if range.start > message.offset_range.start
2533                && range.end < message.offset_range.end.saturating_sub(1)
2534            {
2535                if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
2536                    suffix_start = Some(range.end + 1);
2537                } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
2538                    suffix_start = Some(range.end);
2539                }
2540            }
2541
2542            let version = self.version.clone();
2543            let suffix = if let Some(suffix_start) = suffix_start {
2544                MessageAnchor {
2545                    id: MessageId(self.next_timestamp()),
2546                    start: self.buffer.read(cx).anchor_before(suffix_start),
2547                }
2548            } else {
2549                self.buffer.update(cx, |buffer, cx| {
2550                    buffer.edit([(range.end..range.end, "\n")], None, cx);
2551                });
2552                edited_buffer = true;
2553                MessageAnchor {
2554                    id: MessageId(self.next_timestamp()),
2555                    start: self.buffer.read(cx).anchor_before(range.end + 1),
2556                }
2557            };
2558
2559            let suffix_metadata = MessageMetadata {
2560                role,
2561                status: MessageStatus::Done,
2562                timestamp: suffix.id.0,
2563                cache: None,
2564            };
2565            self.insert_message(suffix.clone(), suffix_metadata.clone(), cx);
2566            self.push_op(
2567                ContextOperation::InsertMessage {
2568                    anchor: suffix.clone(),
2569                    metadata: suffix_metadata,
2570                    version,
2571                },
2572                cx,
2573            );
2574
2575            let new_messages =
2576                if range.start == range.end || range.start == message.offset_range.start {
2577                    (None, Some(suffix))
2578                } else {
2579                    let mut prefix_end = None;
2580                    if range.start > message.offset_range.start
2581                        && range.end < message.offset_range.end - 1
2582                    {
2583                        if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
2584                            prefix_end = Some(range.start + 1);
2585                        } else if self.buffer.read(cx).reversed_chars_at(range.start).next()
2586                            == Some('\n')
2587                        {
2588                            prefix_end = Some(range.start);
2589                        }
2590                    }
2591
2592                    let version = self.version.clone();
2593                    let selection = if let Some(prefix_end) = prefix_end {
2594                        MessageAnchor {
2595                            id: MessageId(self.next_timestamp()),
2596                            start: self.buffer.read(cx).anchor_before(prefix_end),
2597                        }
2598                    } else {
2599                        self.buffer.update(cx, |buffer, cx| {
2600                            buffer.edit([(range.start..range.start, "\n")], None, cx)
2601                        });
2602                        edited_buffer = true;
2603                        MessageAnchor {
2604                            id: MessageId(self.next_timestamp()),
2605                            start: self.buffer.read(cx).anchor_before(range.end + 1),
2606                        }
2607                    };
2608
2609                    let selection_metadata = MessageMetadata {
2610                        role,
2611                        status: MessageStatus::Done,
2612                        timestamp: selection.id.0,
2613                        cache: None,
2614                    };
2615                    self.insert_message(selection.clone(), selection_metadata.clone(), cx);
2616                    self.push_op(
2617                        ContextOperation::InsertMessage {
2618                            anchor: selection.clone(),
2619                            metadata: selection_metadata,
2620                            version,
2621                        },
2622                        cx,
2623                    );
2624
2625                    (Some(selection), Some(suffix))
2626                };
2627
2628            if !edited_buffer {
2629                cx.emit(ContextEvent::MessagesEdited);
2630            }
2631            new_messages
2632        } else {
2633            (None, None)
2634        }
2635    }
2636
2637    fn insert_message(
2638        &mut self,
2639        new_anchor: MessageAnchor,
2640        new_metadata: MessageMetadata,
2641        cx: &mut Context<Self>,
2642    ) {
2643        cx.emit(ContextEvent::MessagesEdited);
2644
2645        self.messages_metadata.insert(new_anchor.id, new_metadata);
2646
2647        let buffer = self.buffer.read(cx);
2648        let insertion_ix = self
2649            .message_anchors
2650            .iter()
2651            .position(|anchor| {
2652                let comparison = new_anchor.start.cmp(&anchor.start, buffer);
2653                comparison.is_lt() || (comparison.is_eq() && new_anchor.id > anchor.id)
2654            })
2655            .unwrap_or(self.message_anchors.len());
2656        self.message_anchors.insert(insertion_ix, new_anchor);
2657    }
2658
2659    pub fn summarize(&mut self, mut replace_old: bool, cx: &mut Context<Self>) {
2660        let Some(model) = LanguageModelRegistry::read_global(cx).default_model() else {
2661            return;
2662        };
2663
2664        if replace_old || (self.message_anchors.len() >= 2 && self.summary.is_pending()) {
2665            if !model.provider.is_authenticated(cx) {
2666                return;
2667            }
2668
2669            let mut request = self.to_completion_request(Some(&model.model), cx);
2670            request.messages.push(LanguageModelRequestMessage {
2671                role: Role::User,
2672                content: vec![
2673                    "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:`"
2674                        .into(),
2675                ],
2676                cache: false,
2677            });
2678
2679            // If there is no summary, it is set with `done: false` so that "Loading Summary…" can
2680            // be displayed.
2681            match self.summary {
2682                ContextSummary::Pending | ContextSummary::Error => {
2683                    self.summary = ContextSummary::Content(ContextSummaryContent {
2684                        text: "".to_string(),
2685                        done: false,
2686                        timestamp: clock::Lamport::default(),
2687                    });
2688                    replace_old = true;
2689                }
2690                ContextSummary::Content(_) => {}
2691            }
2692
2693            self.summary_task = cx.spawn(async move |this, cx| {
2694                let result = async {
2695                    let stream = model.model.stream_completion_text(request, &cx);
2696                    let mut messages = stream.await?;
2697
2698                    let mut replaced = !replace_old;
2699                    while let Some(message) = messages.stream.next().await {
2700                        let text = message?;
2701                        let mut lines = text.lines();
2702                        this.update(cx, |this, cx| {
2703                            let version = this.version.clone();
2704                            let timestamp = this.next_timestamp();
2705                            let summary = this.summary.content_or_set_empty();
2706                            if !replaced && replace_old {
2707                                summary.text.clear();
2708                                replaced = true;
2709                            }
2710                            summary.text.extend(lines.next());
2711                            summary.timestamp = timestamp;
2712                            let operation = ContextOperation::UpdateSummary {
2713                                summary: summary.clone(),
2714                                version,
2715                            };
2716                            this.push_op(operation, cx);
2717                            cx.emit(ContextEvent::SummaryChanged);
2718                            cx.emit(ContextEvent::SummaryGenerated);
2719                        })?;
2720
2721                        // Stop if the LLM generated multiple lines.
2722                        if lines.next().is_some() {
2723                            break;
2724                        }
2725                    }
2726
2727                    this.read_with(cx, |this, _cx| {
2728                        if let Some(summary) = this.summary.content() {
2729                            if summary.text.is_empty() {
2730                                bail!("Model generated an empty summary");
2731                            }
2732                        }
2733                        Ok(())
2734                    })??;
2735
2736                    this.update(cx, |this, cx| {
2737                        let version = this.version.clone();
2738                        let timestamp = this.next_timestamp();
2739                        if let Some(summary) = this.summary.content_as_mut() {
2740                            summary.done = true;
2741                            summary.timestamp = timestamp;
2742                            let operation = ContextOperation::UpdateSummary {
2743                                summary: summary.clone(),
2744                                version,
2745                            };
2746                            this.push_op(operation, cx);
2747                            cx.emit(ContextEvent::SummaryChanged);
2748                            cx.emit(ContextEvent::SummaryGenerated);
2749                        }
2750                    })?;
2751
2752                    anyhow::Ok(())
2753                }
2754                .await;
2755
2756                if let Err(err) = result {
2757                    this.update(cx, |this, cx| {
2758                        this.summary = ContextSummary::Error;
2759                        cx.emit(ContextEvent::SummaryChanged);
2760                    })
2761                    .log_err();
2762                    log::error!("Error generating context summary: {}", err);
2763                }
2764
2765                Some(())
2766            });
2767        }
2768    }
2769
2770    fn message_for_offset(&self, offset: usize, cx: &App) -> Option<Message> {
2771        self.messages_for_offsets([offset], cx).pop()
2772    }
2773
2774    pub fn messages_for_offsets(
2775        &self,
2776        offsets: impl IntoIterator<Item = usize>,
2777        cx: &App,
2778    ) -> Vec<Message> {
2779        let mut result = Vec::new();
2780
2781        let mut messages = self.messages(cx).peekable();
2782        let mut offsets = offsets.into_iter().peekable();
2783        let mut current_message = messages.next();
2784        while let Some(offset) = offsets.next() {
2785            // Locate the message that contains the offset.
2786            while current_message.as_ref().map_or(false, |message| {
2787                !message.offset_range.contains(&offset) && messages.peek().is_some()
2788            }) {
2789                current_message = messages.next();
2790            }
2791            let Some(message) = current_message.as_ref() else {
2792                break;
2793            };
2794
2795            // Skip offsets that are in the same message.
2796            while offsets.peek().map_or(false, |offset| {
2797                message.offset_range.contains(offset) || messages.peek().is_none()
2798            }) {
2799                offsets.next();
2800            }
2801
2802            result.push(message.clone());
2803        }
2804        result
2805    }
2806
2807    fn messages_from_anchors<'a>(
2808        &'a self,
2809        message_anchors: impl Iterator<Item = &'a MessageAnchor> + 'a,
2810        cx: &'a App,
2811    ) -> impl 'a + Iterator<Item = Message> {
2812        let buffer = self.buffer.read(cx);
2813
2814        Self::messages_from_iters(buffer, &self.messages_metadata, message_anchors.enumerate())
2815    }
2816
2817    pub fn messages<'a>(&'a self, cx: &'a App) -> impl 'a + Iterator<Item = Message> {
2818        self.messages_from_anchors(self.message_anchors.iter(), cx)
2819    }
2820
2821    pub fn messages_from_iters<'a>(
2822        buffer: &'a Buffer,
2823        metadata: &'a HashMap<MessageId, MessageMetadata>,
2824        messages: impl Iterator<Item = (usize, &'a MessageAnchor)> + 'a,
2825    ) -> impl 'a + Iterator<Item = Message> {
2826        let mut messages = messages.peekable();
2827
2828        iter::from_fn(move || {
2829            if let Some((start_ix, message_anchor)) = messages.next() {
2830                let metadata = metadata.get(&message_anchor.id)?;
2831
2832                let message_start = message_anchor.start.to_offset(buffer);
2833                let mut message_end = None;
2834                let mut end_ix = start_ix;
2835                while let Some((_, next_message)) = messages.peek() {
2836                    if next_message.start.is_valid(buffer) {
2837                        message_end = Some(next_message.start);
2838                        break;
2839                    } else {
2840                        end_ix += 1;
2841                        messages.next();
2842                    }
2843                }
2844                let message_end_anchor = message_end.unwrap_or(language::Anchor::MAX);
2845                let message_end = message_end_anchor.to_offset(buffer);
2846
2847                return Some(Message {
2848                    index_range: start_ix..end_ix,
2849                    offset_range: message_start..message_end,
2850                    anchor_range: message_anchor.start..message_end_anchor,
2851                    id: message_anchor.id,
2852                    role: metadata.role,
2853                    status: metadata.status.clone(),
2854                    cache: metadata.cache.clone(),
2855                });
2856            }
2857            None
2858        })
2859    }
2860
2861    pub fn save(
2862        &mut self,
2863        debounce: Option<Duration>,
2864        fs: Arc<dyn Fs>,
2865        cx: &mut Context<AssistantContext>,
2866    ) {
2867        if self.replica_id() != ReplicaId::default() {
2868            // Prevent saving a remote context for now.
2869            return;
2870        }
2871
2872        self.pending_save = cx.spawn(async move |this, cx| {
2873            if let Some(debounce) = debounce {
2874                cx.background_executor().timer(debounce).await;
2875            }
2876
2877            let (old_path, summary) = this.read_with(cx, |this, _| {
2878                let path = this.path.clone();
2879                let summary = if let Some(summary) = this.summary.content() {
2880                    if summary.done {
2881                        Some(summary.text.clone())
2882                    } else {
2883                        None
2884                    }
2885                } else {
2886                    None
2887                };
2888                (path, summary)
2889            })?;
2890
2891            if let Some(summary) = summary {
2892                let context = this.read_with(cx, |this, cx| this.serialize(cx))?;
2893                let mut discriminant = 1;
2894                let mut new_path;
2895                loop {
2896                    new_path = contexts_dir().join(&format!(
2897                        "{} - {}.zed.json",
2898                        summary.trim(),
2899                        discriminant
2900                    ));
2901                    if fs.is_file(&new_path).await {
2902                        discriminant += 1;
2903                    } else {
2904                        break;
2905                    }
2906                }
2907
2908                fs.create_dir(contexts_dir().as_ref()).await?;
2909
2910                // rename before write ensures that only one file exists
2911                if let Some(old_path) = old_path.as_ref() {
2912                    if new_path.as_path() != old_path.as_ref() {
2913                        fs.rename(
2914                            &old_path,
2915                            &new_path,
2916                            RenameOptions {
2917                                overwrite: true,
2918                                ignore_if_exists: true,
2919                            },
2920                        )
2921                        .await?;
2922                    }
2923                }
2924
2925                // update path before write in case it fails
2926                this.update(cx, {
2927                    let new_path: Arc<Path> = new_path.clone().into();
2928                    move |this, cx| {
2929                        this.path = Some(new_path.clone());
2930                        cx.emit(ContextEvent::PathChanged { old_path, new_path });
2931                    }
2932                })
2933                .ok();
2934
2935                fs.atomic_write(new_path, serde_json::to_string(&context).unwrap())
2936                    .await?;
2937            }
2938
2939            Ok(())
2940        });
2941    }
2942
2943    pub fn set_custom_summary(&mut self, custom_summary: String, cx: &mut Context<Self>) {
2944        let timestamp = self.next_timestamp();
2945        let summary = self.summary.content_or_set_empty();
2946        summary.timestamp = timestamp;
2947        summary.done = true;
2948        summary.text = custom_summary;
2949        cx.emit(ContextEvent::SummaryChanged);
2950    }
2951}
2952
2953#[derive(Debug, Default)]
2954pub struct ContextVersion {
2955    context: clock::Global,
2956    buffer: clock::Global,
2957}
2958
2959impl ContextVersion {
2960    pub fn from_proto(proto: &proto::ContextVersion) -> Self {
2961        Self {
2962            context: language::proto::deserialize_version(&proto.context_version),
2963            buffer: language::proto::deserialize_version(&proto.buffer_version),
2964        }
2965    }
2966
2967    pub fn to_proto(&self, context_id: ContextId) -> proto::ContextVersion {
2968        proto::ContextVersion {
2969            context_id: context_id.to_proto(),
2970            context_version: language::proto::serialize_version(&self.context),
2971            buffer_version: language::proto::serialize_version(&self.buffer),
2972        }
2973    }
2974}
2975
2976#[derive(Debug, Clone)]
2977pub struct ParsedSlashCommand {
2978    pub name: String,
2979    pub arguments: SmallVec<[String; 3]>,
2980    pub status: PendingSlashCommandStatus,
2981    pub source_range: Range<language::Anchor>,
2982}
2983
2984#[derive(Debug)]
2985pub struct InvokedSlashCommand {
2986    pub name: SharedString,
2987    pub range: Range<language::Anchor>,
2988    pub run_commands_in_ranges: Vec<Range<language::Anchor>>,
2989    pub status: InvokedSlashCommandStatus,
2990    pub transaction: Option<language::TransactionId>,
2991    timestamp: clock::Lamport,
2992}
2993
2994#[derive(Debug)]
2995pub enum InvokedSlashCommandStatus {
2996    Running(Task<()>),
2997    Error(SharedString),
2998    Finished,
2999}
3000
3001#[derive(Debug, Clone)]
3002pub enum PendingSlashCommandStatus {
3003    Idle,
3004    Running { _task: Shared<Task<()>> },
3005    Error(String),
3006}
3007
3008#[derive(Debug, Clone)]
3009pub struct PendingToolUse {
3010    pub id: LanguageModelToolUseId,
3011    pub name: String,
3012    pub input: serde_json::Value,
3013    pub status: PendingToolUseStatus,
3014    pub source_range: Range<language::Anchor>,
3015}
3016
3017#[derive(Debug, Clone)]
3018pub enum PendingToolUseStatus {
3019    Idle,
3020    Running { _task: Shared<Task<()>> },
3021    Error(String),
3022}
3023
3024impl PendingToolUseStatus {
3025    pub fn is_idle(&self) -> bool {
3026        matches!(self, PendingToolUseStatus::Idle)
3027    }
3028}
3029
3030#[derive(Serialize, Deserialize)]
3031pub struct SavedMessage {
3032    pub id: MessageId,
3033    pub start: usize,
3034    pub metadata: MessageMetadata,
3035}
3036
3037#[derive(Serialize, Deserialize)]
3038pub struct SavedContext {
3039    pub id: Option<ContextId>,
3040    pub zed: String,
3041    pub version: String,
3042    pub text: String,
3043    pub messages: Vec<SavedMessage>,
3044    pub summary: String,
3045    pub slash_command_output_sections:
3046        Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
3047    #[serde(default)]
3048    pub thought_process_output_sections: Vec<ThoughtProcessOutputSection<usize>>,
3049}
3050
3051impl SavedContext {
3052    pub const VERSION: &'static str = "0.4.0";
3053
3054    pub fn from_json(json: &str) -> Result<Self> {
3055        let saved_context_json = serde_json::from_str::<serde_json::Value>(json)?;
3056        match saved_context_json
3057            .get("version")
3058            .context("version not found")?
3059        {
3060            serde_json::Value::String(version) => match version.as_str() {
3061                SavedContext::VERSION => {
3062                    Ok(serde_json::from_value::<SavedContext>(saved_context_json)?)
3063                }
3064                SavedContextV0_3_0::VERSION => {
3065                    let saved_context =
3066                        serde_json::from_value::<SavedContextV0_3_0>(saved_context_json)?;
3067                    Ok(saved_context.upgrade())
3068                }
3069                SavedContextV0_2_0::VERSION => {
3070                    let saved_context =
3071                        serde_json::from_value::<SavedContextV0_2_0>(saved_context_json)?;
3072                    Ok(saved_context.upgrade())
3073                }
3074                SavedContextV0_1_0::VERSION => {
3075                    let saved_context =
3076                        serde_json::from_value::<SavedContextV0_1_0>(saved_context_json)?;
3077                    Ok(saved_context.upgrade())
3078                }
3079                _ => anyhow::bail!("unrecognized saved context version: {version:?}"),
3080            },
3081            _ => anyhow::bail!("version not found on saved context"),
3082        }
3083    }
3084
3085    fn into_ops(
3086        self,
3087        buffer: &Entity<Buffer>,
3088        cx: &mut Context<AssistantContext>,
3089    ) -> Vec<ContextOperation> {
3090        let mut operations = Vec::new();
3091        let mut version = clock::Global::new();
3092        let mut next_timestamp = clock::Lamport::new(ReplicaId::default());
3093
3094        let mut first_message_metadata = None;
3095        for message in self.messages {
3096            if message.id == MessageId(clock::Lamport::default()) {
3097                first_message_metadata = Some(message.metadata);
3098            } else {
3099                operations.push(ContextOperation::InsertMessage {
3100                    anchor: MessageAnchor {
3101                        id: message.id,
3102                        start: buffer.read(cx).anchor_before(message.start),
3103                    },
3104                    metadata: MessageMetadata {
3105                        role: message.metadata.role,
3106                        status: message.metadata.status,
3107                        timestamp: message.metadata.timestamp,
3108                        cache: None,
3109                    },
3110                    version: version.clone(),
3111                });
3112                version.observe(message.id.0);
3113                next_timestamp.observe(message.id.0);
3114            }
3115        }
3116
3117        if let Some(metadata) = first_message_metadata {
3118            let timestamp = next_timestamp.tick();
3119            operations.push(ContextOperation::UpdateMessage {
3120                message_id: MessageId(clock::Lamport::default()),
3121                metadata: MessageMetadata {
3122                    role: metadata.role,
3123                    status: metadata.status,
3124                    timestamp,
3125                    cache: None,
3126                },
3127                version: version.clone(),
3128            });
3129            version.observe(timestamp);
3130        }
3131
3132        let buffer = buffer.read(cx);
3133        for section in self.slash_command_output_sections {
3134            let timestamp = next_timestamp.tick();
3135            operations.push(ContextOperation::SlashCommandOutputSectionAdded {
3136                timestamp,
3137                section: SlashCommandOutputSection {
3138                    range: buffer.anchor_after(section.range.start)
3139                        ..buffer.anchor_before(section.range.end),
3140                    icon: section.icon,
3141                    label: section.label,
3142                    metadata: section.metadata,
3143                },
3144                version: version.clone(),
3145            });
3146
3147            version.observe(timestamp);
3148        }
3149
3150        for section in self.thought_process_output_sections {
3151            let timestamp = next_timestamp.tick();
3152            operations.push(ContextOperation::ThoughtProcessOutputSectionAdded {
3153                timestamp,
3154                section: ThoughtProcessOutputSection {
3155                    range: buffer.anchor_after(section.range.start)
3156                        ..buffer.anchor_before(section.range.end),
3157                },
3158                version: version.clone(),
3159            });
3160
3161            version.observe(timestamp);
3162        }
3163
3164        let timestamp = next_timestamp.tick();
3165        operations.push(ContextOperation::UpdateSummary {
3166            summary: ContextSummaryContent {
3167                text: self.summary,
3168                done: true,
3169                timestamp,
3170            },
3171            version: version.clone(),
3172        });
3173        version.observe(timestamp);
3174
3175        operations
3176    }
3177}
3178
3179#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
3180struct SavedMessageIdPreV0_4_0(usize);
3181
3182#[derive(Serialize, Deserialize)]
3183struct SavedMessagePreV0_4_0 {
3184    id: SavedMessageIdPreV0_4_0,
3185    start: usize,
3186}
3187
3188#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
3189struct SavedMessageMetadataPreV0_4_0 {
3190    role: Role,
3191    status: MessageStatus,
3192}
3193
3194#[derive(Serialize, Deserialize)]
3195struct SavedContextV0_3_0 {
3196    id: Option<ContextId>,
3197    zed: String,
3198    version: String,
3199    text: String,
3200    messages: Vec<SavedMessagePreV0_4_0>,
3201    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3202    summary: String,
3203    slash_command_output_sections: Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
3204}
3205
3206impl SavedContextV0_3_0 {
3207    const VERSION: &'static str = "0.3.0";
3208
3209    fn upgrade(self) -> SavedContext {
3210        SavedContext {
3211            id: self.id,
3212            zed: self.zed,
3213            version: SavedContext::VERSION.into(),
3214            text: self.text,
3215            messages: self
3216                .messages
3217                .into_iter()
3218                .filter_map(|message| {
3219                    let metadata = self.message_metadata.get(&message.id)?;
3220                    let timestamp = clock::Lamport {
3221                        replica_id: ReplicaId::default(),
3222                        value: message.id.0 as u32,
3223                    };
3224                    Some(SavedMessage {
3225                        id: MessageId(timestamp),
3226                        start: message.start,
3227                        metadata: MessageMetadata {
3228                            role: metadata.role,
3229                            status: metadata.status.clone(),
3230                            timestamp,
3231                            cache: None,
3232                        },
3233                    })
3234                })
3235                .collect(),
3236            summary: self.summary,
3237            slash_command_output_sections: self.slash_command_output_sections,
3238            thought_process_output_sections: Vec::new(),
3239        }
3240    }
3241}
3242
3243#[derive(Serialize, Deserialize)]
3244struct SavedContextV0_2_0 {
3245    id: Option<ContextId>,
3246    zed: String,
3247    version: String,
3248    text: String,
3249    messages: Vec<SavedMessagePreV0_4_0>,
3250    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3251    summary: String,
3252}
3253
3254impl SavedContextV0_2_0 {
3255    const VERSION: &'static str = "0.2.0";
3256
3257    fn upgrade(self) -> SavedContext {
3258        SavedContextV0_3_0 {
3259            id: self.id,
3260            zed: self.zed,
3261            version: SavedContextV0_3_0::VERSION.to_string(),
3262            text: self.text,
3263            messages: self.messages,
3264            message_metadata: self.message_metadata,
3265            summary: self.summary,
3266            slash_command_output_sections: Vec::new(),
3267        }
3268        .upgrade()
3269    }
3270}
3271
3272#[derive(Serialize, Deserialize)]
3273struct SavedContextV0_1_0 {
3274    id: Option<ContextId>,
3275    zed: String,
3276    version: String,
3277    text: String,
3278    messages: Vec<SavedMessagePreV0_4_0>,
3279    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3280    summary: String,
3281    api_url: Option<String>,
3282    model: OpenAiModel,
3283}
3284
3285impl SavedContextV0_1_0 {
3286    const VERSION: &'static str = "0.1.0";
3287
3288    fn upgrade(self) -> SavedContext {
3289        SavedContextV0_2_0 {
3290            id: self.id,
3291            zed: self.zed,
3292            version: SavedContextV0_2_0::VERSION.to_string(),
3293            text: self.text,
3294            messages: self.messages,
3295            message_metadata: self.message_metadata,
3296            summary: self.summary,
3297        }
3298        .upgrade()
3299    }
3300}
3301
3302#[derive(Debug, Clone)]
3303pub struct SavedContextMetadata {
3304    pub title: SharedString,
3305    pub path: Arc<Path>,
3306    pub mtime: chrono::DateTime<chrono::Local>,
3307}