Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
ZIm/crates/assistant_context/src/assistant_context.rs
Bennet Bo Fenner 41fe2a2ab4
agent: Disable thinking when using inline assistant/edit file tool (#34141) 
This introduces a new field `thinking_allowed` on `LanguageModelRequest`
which lets us control whether thinking should be enabled if the model
supports it.
We permit thinking in the Inline Assistant, Edit File tool and the Git
Commit message generator, this should make generation faster when using
a thinking model, e.g. `claude-sonnet-4-thinking`

Release Notes:

- N/A
2025-07-09 18:05:39 +00:00

3315 lines
123 KiB
Rust
Raw Blame History

#[cfg(test)]
mod assistant_context_tests;
mod context_store;
use agent_settings::AgentSettings;
use anyhow::{Context as _, Result, bail};
use assistant_slash_command::{
SlashCommandContent, SlashCommandEvent, SlashCommandLine, SlashCommandOutputSection,
SlashCommandResult, SlashCommandWorkingSet,
};
use assistant_slash_commands::FileCommandMetadata;
use client::{self, Client, proto, telemetry::Telemetry};
use clock::ReplicaId;
use collections::{HashMap, HashSet};
use fs::{Fs, RenameOptions};
use futures::{FutureExt, StreamExt, future::Shared};
use gpui::{
App, AppContext as _, Context, Entity, EventEmitter, RenderImage, SharedString, Subscription,
Task,
};
use language::{AnchorRangeExt, Bias, Buffer, LanguageRegistry, OffsetRangeExt, Point, ToOffset};
use language_model::{
LanguageModel, LanguageModelCacheConfiguration, LanguageModelCompletionEvent,
LanguageModelImage, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
LanguageModelToolUseId, MessageContent, PaymentRequiredError, Role, StopReason,
report_assistant_event,
};
use open_ai::Model as OpenAiModel;
use paths::contexts_dir;
use project::Project;
use prompt_store::PromptBuilder;
use serde::{Deserialize, Serialize};
use settings::Settings;
use smallvec::SmallVec;
use std::{
cmp::{Ordering, max},
fmt::{Debug, Write as _},
iter, mem,
ops::Range,
path::Path,
sync::Arc,
time::{Duration, Instant},
};
use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};
use text::{BufferSnapshot, ToPoint};
use ui::IconName;
use util::{ResultExt, TryFutureExt, post_inc};
use uuid::Uuid;
use zed_llm_client::CompletionIntent;
pub use crate::context_store::*;
pub fn init(client: Arc<Client>, _: &mut App) {
context_store::init(&client.into());
}
#[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ContextId(String);
impl ContextId {
pub fn new() -> Self {
Self(Uuid::new_v4().to_string())
}
pub fn from_proto(id: String) -> Self {
Self(id)
}
pub fn to_proto(&self) -> String {
self.0.clone()
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct MessageId(pub clock::Lamport);
impl MessageId {
pub fn as_u64(self) -> u64 {
self.0.as_u64()
}
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum MessageStatus {
Pending,
Done,
Error(SharedString),
Canceled,
}
impl MessageStatus {
pub fn from_proto(status: proto::ContextMessageStatus) -> MessageStatus {
match status.variant {
Some(proto::context_message_status::Variant::Pending(_)) => MessageStatus::Pending,
Some(proto::context_message_status::Variant::Done(_)) => MessageStatus::Done,
Some(proto::context_message_status::Variant::Error(error)) => {
MessageStatus::Error(error.message.into())
}
Some(proto::context_message_status::Variant::Canceled(_)) => MessageStatus::Canceled,
None => MessageStatus::Pending,
}
}
pub fn to_proto(&self) -> proto::ContextMessageStatus {
match self {
MessageStatus::Pending => proto::ContextMessageStatus {
variant: Some(proto::context_message_status::Variant::Pending(
proto::context_message_status::Pending {},
)),
},
MessageStatus::Done => proto::ContextMessageStatus {
variant: Some(proto::context_message_status::Variant::Done(
proto::context_message_status::Done {},
)),
},
MessageStatus::Error(message) => proto::ContextMessageStatus {
variant: Some(proto::context_message_status::Variant::Error(
proto::context_message_status::Error {
message: message.to_string(),
},
)),
},
MessageStatus::Canceled => proto::ContextMessageStatus {
variant: Some(proto::context_message_status::Variant::Canceled(
proto::context_message_status::Canceled {},
)),
},
}
}
}
#[derive(Clone, Debug)]
pub enum ContextOperation {
InsertMessage {
anchor: MessageAnchor,
metadata: MessageMetadata,
version: clock::Global,
},
UpdateMessage {
message_id: MessageId,
metadata: MessageMetadata,
version: clock::Global,
},
UpdateSummary {
summary: ContextSummaryContent,
version: clock::Global,
},
SlashCommandStarted {
id: InvokedSlashCommandId,
output_range: Range<language::Anchor>,
name: String,
version: clock::Global,
},
SlashCommandFinished {
id: InvokedSlashCommandId,
timestamp: clock::Lamport,
error_message: Option<String>,
version: clock::Global,
},
SlashCommandOutputSectionAdded {
timestamp: clock::Lamport,
section: SlashCommandOutputSection<language::Anchor>,
version: clock::Global,
},
ThoughtProcessOutputSectionAdded {
timestamp: clock::Lamport,
section: ThoughtProcessOutputSection<language::Anchor>,
version: clock::Global,
},
BufferOperation(language::Operation),
}
impl ContextOperation {
pub fn from_proto(op: proto::ContextOperation) -> Result<Self> {
match op.variant.context("invalid variant")? {
proto::context_operation::Variant::InsertMessage(insert) => {
let message = insert.message.context("invalid message")?;
let id = MessageId(language::proto::deserialize_timestamp(
message.id.context("invalid id")?,
));
Ok(Self::InsertMessage {
anchor: MessageAnchor {
id,
start: language::proto::deserialize_anchor(
message.start.context("invalid anchor")?,
)
.context("invalid anchor")?,
},
metadata: MessageMetadata {
role: Role::from_proto(message.role),
status: MessageStatus::from_proto(
message.status.context("invalid status")?,
),
timestamp: id.0,
cache: None,
},
version: language::proto::deserialize_version(&insert.version),
})
}
proto::context_operation::Variant::UpdateMessage(update) => Ok(Self::UpdateMessage {
message_id: MessageId(language::proto::deserialize_timestamp(
update.message_id.context("invalid message id")?,
)),
metadata: MessageMetadata {
role: Role::from_proto(update.role),
status: MessageStatus::from_proto(update.status.context("invalid status")?),
timestamp: language::proto::deserialize_timestamp(
update.timestamp.context("invalid timestamp")?,
),
cache: None,
},
version: language::proto::deserialize_version(&update.version),
}),
proto::context_operation::Variant::UpdateSummary(update) => Ok(Self::UpdateSummary {
summary: ContextSummaryContent {
text: update.summary,
done: update.done,
timestamp: language::proto::deserialize_timestamp(
update.timestamp.context("invalid timestamp")?,
),
},
version: language::proto::deserialize_version(&update.version),
}),
proto::context_operation::Variant::SlashCommandStarted(message) => {
Ok(Self::SlashCommandStarted {
id: InvokedSlashCommandId(language::proto::deserialize_timestamp(
message.id.context("invalid id")?,
)),
output_range: language::proto::deserialize_anchor_range(
message.output_range.context("invalid range")?,
)?,
name: message.name,
version: language::proto::deserialize_version(&message.version),
})
}
proto::context_operation::Variant::SlashCommandOutputSectionAdded(message) => {
let section = message.section.context("missing section")?;
Ok(Self::SlashCommandOutputSectionAdded {
timestamp: language::proto::deserialize_timestamp(
message.timestamp.context("missing timestamp")?,
),
section: SlashCommandOutputSection {
range: language::proto::deserialize_anchor_range(
section.range.context("invalid range")?,
)?,
icon: section.icon_name.parse()?,
label: section.label.into(),
metadata: section
.metadata
.and_then(|metadata| serde_json::from_str(&metadata).log_err()),
},
version: language::proto::deserialize_version(&message.version),
})
}
proto::context_operation::Variant::SlashCommandCompleted(message) => {
Ok(Self::SlashCommandFinished {
id: InvokedSlashCommandId(language::proto::deserialize_timestamp(
message.id.context("invalid id")?,
)),
timestamp: language::proto::deserialize_timestamp(
message.timestamp.context("missing timestamp")?,
),
error_message: message.error_message,
version: language::proto::deserialize_version(&message.version),
})
}
proto::context_operation::Variant::ThoughtProcessOutputSectionAdded(message) => {
let section = message.section.context("missing section")?;
Ok(Self::ThoughtProcessOutputSectionAdded {
timestamp: language::proto::deserialize_timestamp(
message.timestamp.context("missing timestamp")?,
),
section: ThoughtProcessOutputSection {
range: language::proto::deserialize_anchor_range(
section.range.context("invalid range")?,
)?,
},
version: language::proto::deserialize_version(&message.version),
})
}
proto::context_operation::Variant::BufferOperation(op) => Ok(Self::BufferOperation(
language::proto::deserialize_operation(
op.operation.context("invalid buffer operation")?,
)?,
)),
}
}
pub fn to_proto(&self) -> proto::ContextOperation {
match self {
Self::InsertMessage {
anchor,
metadata,
version,
} => proto::ContextOperation {
variant: Some(proto::context_operation::Variant::InsertMessage(
proto::context_operation::InsertMessage {
message: Some(proto::ContextMessage {
id: Some(language::proto::serialize_timestamp(anchor.id.0)),
start: Some(language::proto::serialize_anchor(&anchor.start)),
role: metadata.role.to_proto() as i32,
status: Some(metadata.status.to_proto()),
}),
version: language::proto::serialize_version(version),
},
)),
},
Self::UpdateMessage {
message_id,
metadata,
version,
} => proto::ContextOperation {
variant: Some(proto::context_operation::Variant::UpdateMessage(
proto::context_operation::UpdateMessage {
message_id: Some(language::proto::serialize_timestamp(message_id.0)),
role: metadata.role.to_proto() as i32,
status: Some(metadata.status.to_proto()),
timestamp: Some(language::proto::serialize_timestamp(metadata.timestamp)),
version: language::proto::serialize_version(version),
},
)),
},
Self::UpdateSummary { summary, version } => proto::ContextOperation {
variant: Some(proto::context_operation::Variant::UpdateSummary(
proto::context_operation::UpdateSummary {
summary: summary.text.clone(),
done: summary.done,
timestamp: Some(language::proto::serialize_timestamp(summary.timestamp)),
version: language::proto::serialize_version(version),
},
)),
},
Self::SlashCommandStarted {
id,
output_range,
name,
version,
} => proto::ContextOperation {
variant: Some(proto::context_operation::Variant::SlashCommandStarted(
proto::context_operation::SlashCommandStarted {
id: Some(language::proto::serialize_timestamp(id.0)),
output_range: Some(language::proto::serialize_anchor_range(
output_range.clone(),
)),
name: name.clone(),
version: language::proto::serialize_version(version),
},
)),
},
Self::SlashCommandOutputSectionAdded {
timestamp,
section,
version,
} => proto::ContextOperation {
variant: Some(
proto::context_operation::Variant::SlashCommandOutputSectionAdded(
proto::context_operation::SlashCommandOutputSectionAdded {
timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
section: Some({
let icon_name: &'static str = section.icon.into();
proto::SlashCommandOutputSection {
range: Some(language::proto::serialize_anchor_range(
section.range.clone(),
)),
icon_name: icon_name.to_string(),
label: section.label.to_string(),
metadata: section.metadata.as_ref().and_then(|metadata| {
serde_json::to_string(metadata).log_err()
}),
}
}),
version: language::proto::serialize_version(version),
},
),
),
},
Self::SlashCommandFinished {
id,
timestamp,
error_message,
version,
} => proto::ContextOperation {
variant: Some(proto::context_operation::Variant::SlashCommandCompleted(
proto::context_operation::SlashCommandCompleted {
id: Some(language::proto::serialize_timestamp(id.0)),
timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
error_message: error_message.clone(),
version: language::proto::serialize_version(version),
},
)),
},
Self::ThoughtProcessOutputSectionAdded {
timestamp,
section,
version,
} => proto::ContextOperation {
variant: Some(
proto::context_operation::Variant::ThoughtProcessOutputSectionAdded(
proto::context_operation::ThoughtProcessOutputSectionAdded {
timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
section: Some({
proto::ThoughtProcessOutputSection {
range: Some(language::proto::serialize_anchor_range(
section.range.clone(),
)),
}
}),
version: language::proto::serialize_version(version),
},
),
),
},
Self::BufferOperation(operation) => proto::ContextOperation {
variant: Some(proto::context_operation::Variant::BufferOperation(
proto::context_operation::BufferOperation {
operation: Some(language::proto::serialize_operation(operation)),
},
)),
},
}
}
fn timestamp(&self) -> clock::Lamport {
match self {
Self::InsertMessage { anchor, .. } => anchor.id.0,
Self::UpdateMessage { metadata, .. } => metadata.timestamp,
Self::UpdateSummary { summary, .. } => summary.timestamp,
Self::SlashCommandStarted { id, .. } => id.0,
Self::SlashCommandOutputSectionAdded { timestamp, .. }
| Self::SlashCommandFinished { timestamp, .. }
| Self::ThoughtProcessOutputSectionAdded { timestamp, .. } => *timestamp,
Self::BufferOperation(_) => {
panic!("reading the timestamp of a buffer operation is not supported")
}
}
}
/// Returns the current version of the context operation.
pub fn version(&self) -> &clock::Global {
match self {
Self::InsertMessage { version, .. }
| Self::UpdateMessage { version, .. }
| Self::UpdateSummary { version, .. }
| Self::SlashCommandStarted { version, .. }
| Self::SlashCommandOutputSectionAdded { version, .. }
| Self::SlashCommandFinished { version, .. }
| Self::ThoughtProcessOutputSectionAdded { version, .. } => version,
Self::BufferOperation(_) => {
panic!("reading the version of a buffer operation is not supported")
}
}
}
}
#[derive(Debug, Clone)]
pub enum ContextEvent {
ShowAssistError(SharedString),
ShowPaymentRequiredError,
MessagesEdited,
SummaryChanged,
SummaryGenerated,
PathChanged {
old_path: Option<Arc<Path>>,
new_path: Arc<Path>,
},
StreamedCompletion,
StartedThoughtProcess(Range<language::Anchor>),
EndedThoughtProcess(language::Anchor),
InvokedSlashCommandChanged {
command_id: InvokedSlashCommandId,
},
ParsedSlashCommandsUpdated {
removed: Vec<Range<language::Anchor>>,
updated: Vec<ParsedSlashCommand>,
},
SlashCommandOutputSectionAdded {
section: SlashCommandOutputSection<language::Anchor>,
},
Operation(ContextOperation),
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ContextSummary {
Pending,
Content(ContextSummaryContent),
Error,
}
#[derive(Default, Clone, Debug, Eq, PartialEq)]
pub struct ContextSummaryContent {
pub text: String,
pub done: bool,
pub timestamp: clock::Lamport,
}
impl ContextSummary {
pub const DEFAULT: &str = "New Text Thread";
pub fn or_default(&self) -> SharedString {
self.unwrap_or(Self::DEFAULT)
}
pub fn unwrap_or(&self, message: impl Into<SharedString>) -> SharedString {
self.content()
.map_or_else(|| message.into(), |content| content.text.clone().into())
}
pub fn content(&self) -> Option<&ContextSummaryContent> {
match self {
ContextSummary::Content(content) => Some(content),
ContextSummary::Pending | ContextSummary::Error => None,
}
}
fn content_as_mut(&mut self) -> Option<&mut ContextSummaryContent> {
match self {
ContextSummary::Content(content) => Some(content),
ContextSummary::Pending | ContextSummary::Error => None,
}
}
fn content_or_set_empty(&mut self) -> &mut ContextSummaryContent {
match self {
ContextSummary::Content(content) => content,
ContextSummary::Pending | ContextSummary::Error => {
let content = ContextSummaryContent::default();
*self = ContextSummary::Content(content);
self.content_as_mut().unwrap()
}
}
}
pub fn is_pending(&self) -> bool {
matches!(self, ContextSummary::Pending)
}
fn timestamp(&self) -> Option<clock::Lamport> {
match self {
ContextSummary::Content(content) => Some(content.timestamp),
ContextSummary::Pending | ContextSummary::Error => None,
}
}
}
impl PartialOrd for ContextSummary {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
self.timestamp().partial_cmp(&other.timestamp())
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MessageAnchor {
pub id: MessageId,
pub start: language::Anchor,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum CacheStatus {
Pending,
Cached,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MessageCacheMetadata {
pub is_anchor: bool,
pub is_final_anchor: bool,
pub status: CacheStatus,
pub cached_at: clock::Global,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct MessageMetadata {
pub role: Role,
pub status: MessageStatus,
pub timestamp: clock::Lamport,
#[serde(skip)]
pub cache: Option<MessageCacheMetadata>,
}
impl From<&Message> for MessageMetadata {
fn from(message: &Message) -> Self {
Self {
role: message.role,
status: message.status.clone(),
timestamp: message.id.0,
cache: message.cache.clone(),
}
}
}
impl MessageMetadata {
pub fn is_cache_valid(&self, buffer: &BufferSnapshot, range: &Range<usize>) -> bool {
let result = match &self.cache {
Some(MessageCacheMetadata { cached_at, .. }) => !buffer.has_edits_since_in_range(
&cached_at,
Range {
start: buffer.anchor_at(range.start, Bias::Right),
end: buffer.anchor_at(range.end, Bias::Left),
},
),
_ => false,
};
result
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ThoughtProcessOutputSection<T> {
pub range: Range<T>,
}
impl ThoughtProcessOutputSection<language::Anchor> {
pub fn is_valid(&self, buffer: &language::TextBuffer) -> bool {
self.range.start.is_valid(buffer) && !self.range.to_offset(buffer).is_empty()
}
}
#[derive(Clone, Debug)]
pub struct Message {
pub offset_range: Range<usize>,
pub index_range: Range<usize>,
pub anchor_range: Range<language::Anchor>,
pub id: MessageId,
pub role: Role,
pub status: MessageStatus,
pub cache: Option<MessageCacheMetadata>,
}
#[derive(Debug, Clone)]
pub enum Content {
Image {
anchor: language::Anchor,
image_id: u64,
render_image: Arc<RenderImage>,
image: Shared<Task<Option<LanguageModelImage>>>,
},
}
impl Content {
fn range(&self) -> Range<language::Anchor> {
match self {
Self::Image { anchor, .. } => *anchor..*anchor,
}
}
fn cmp(&self, other: &Self, buffer: &BufferSnapshot) -> Ordering {
let self_range = self.range();
let other_range = other.range();
if self_range.end.cmp(&other_range.start, buffer).is_lt() {
Ordering::Less
} else if self_range.start.cmp(&other_range.end, buffer).is_gt() {
Ordering::Greater
} else {
Ordering::Equal
}
}
}
struct PendingCompletion {
id: usize,
assistant_message_id: MessageId,
_task: Task<()>,
}
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
pub struct InvokedSlashCommandId(clock::Lamport);
pub struct AssistantContext {
id: ContextId,
timestamp: clock::Lamport,
version: clock::Global,
pending_ops: Vec<ContextOperation>,
operations: Vec<ContextOperation>,
buffer: Entity<Buffer>,
parsed_slash_commands: Vec<ParsedSlashCommand>,
invoked_slash_commands: HashMap<InvokedSlashCommandId, InvokedSlashCommand>,
edits_since_last_parse: language::Subscription,
slash_commands: Arc<SlashCommandWorkingSet>,
slash_command_output_sections: Vec<SlashCommandOutputSection<language::Anchor>>,
thought_process_output_sections: Vec<ThoughtProcessOutputSection<language::Anchor>>,
message_anchors: Vec<MessageAnchor>,
contents: Vec<Content>,
messages_metadata: HashMap<MessageId, MessageMetadata>,
summary: ContextSummary,
summary_task: Task<Option<()>>,
completion_count: usize,
pending_completions: Vec<PendingCompletion>,
token_count: Option<u64>,
pending_token_count: Task<Option<()>>,
pending_save: Task<Result<()>>,
pending_cache_warming_task: Task<Option<()>>,
path: Option<Arc<Path>>,
_subscriptions: Vec<Subscription>,
telemetry: Option<Arc<Telemetry>>,
language_registry: Arc<LanguageRegistry>,
project: Option<Entity<Project>>,
prompt_builder: Arc<PromptBuilder>,
completion_mode: agent_settings::CompletionMode,
}
trait ContextAnnotation {
fn range(&self) -> &Range<language::Anchor>;
}
impl ContextAnnotation for ParsedSlashCommand {
fn range(&self) -> &Range<language::Anchor> {
&self.source_range
}
}
impl EventEmitter<ContextEvent> for AssistantContext {}
impl AssistantContext {
pub fn local(
language_registry: Arc<LanguageRegistry>,
project: Option<Entity<Project>>,
telemetry: Option<Arc<Telemetry>>,
prompt_builder: Arc<PromptBuilder>,
slash_commands: Arc<SlashCommandWorkingSet>,
cx: &mut Context<Self>,
) -> Self {
Self::new(
ContextId::new(),
ReplicaId::default(),
language::Capability::ReadWrite,
language_registry,
prompt_builder,
slash_commands,
project,
telemetry,
cx,
)
}
pub fn completion_mode(&self) -> agent_settings::CompletionMode {
self.completion_mode
}
pub fn set_completion_mode(&mut self, completion_mode: agent_settings::CompletionMode) {
self.completion_mode = completion_mode;
}
pub fn new(
id: ContextId,
replica_id: ReplicaId,
capability: language::Capability,
language_registry: Arc<LanguageRegistry>,
prompt_builder: Arc<PromptBuilder>,
slash_commands: Arc<SlashCommandWorkingSet>,
project: Option<Entity<Project>>,
telemetry: Option<Arc<Telemetry>>,
cx: &mut Context<Self>,
) -> Self {
let buffer = cx.new(|_cx| {
let buffer = Buffer::remote(
language::BufferId::new(1).unwrap(),
replica_id,
capability,
"",
);
buffer.set_language_registry(language_registry.clone());
buffer
});
let edits_since_last_slash_command_parse =
buffer.update(cx, |buffer, _| buffer.subscribe());
let mut this = Self {
id,
timestamp: clock::Lamport::new(replica_id),
version: clock::Global::new(),
pending_ops: Vec::new(),
operations: Vec::new(),
message_anchors: Default::default(),
contents: Default::default(),
messages_metadata: Default::default(),
parsed_slash_commands: Vec::new(),
invoked_slash_commands: HashMap::default(),
slash_command_output_sections: Vec::new(),
thought_process_output_sections: Vec::new(),
edits_since_last_parse: edits_since_last_slash_command_parse,
summary: ContextSummary::Pending,
summary_task: Task::ready(None),
completion_count: Default::default(),
pending_completions: Default::default(),
token_count: None,
pending_token_count: Task::ready(None),
pending_cache_warming_task: Task::ready(None),
_subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
pending_save: Task::ready(Ok(())),
completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,
path: None,
buffer,
telemetry,
project,
language_registry,
slash_commands,
prompt_builder,
};
let first_message_id = MessageId(clock::Lamport {
replica_id: 0,
value: 0,
});
let message = MessageAnchor {
id: first_message_id,
start: language::Anchor::MIN,
};
this.messages_metadata.insert(
first_message_id,
MessageMetadata {
role: Role::User,
status: MessageStatus::Done,
timestamp: first_message_id.0,
cache: None,
},
);
this.message_anchors.push(message);
this.set_language(cx);
this.count_remaining_tokens(cx);
this
}
pub(crate) fn serialize(&self, cx: &App) -> SavedContext {
let buffer = self.buffer.read(cx);
SavedContext {
id: Some(self.id.clone()),
zed: "context".into(),
version: SavedContext::VERSION.into(),
text: buffer.text(),
messages: self
.messages(cx)
.map(|message| SavedMessage {
id: message.id,
start: message.offset_range.start,
metadata: self.messages_metadata[&message.id].clone(),
})
.collect(),
summary: self
.summary
.content()
.map(|summary| summary.text.clone())
.unwrap_or_default(),
slash_command_output_sections: self
.slash_command_output_sections
.iter()
.filter_map(|section| {
if section.is_valid(buffer) {
let range = section.range.to_offset(buffer);
Some(assistant_slash_command::SlashCommandOutputSection {
range,
icon: section.icon,
label: section.label.clone(),
metadata: section.metadata.clone(),
})
} else {
None
}
})
.collect(),
thought_process_output_sections: self
.thought_process_output_sections
.iter()
.filter_map(|section| {
if section.is_valid(buffer) {
let range = section.range.to_offset(buffer);
Some(ThoughtProcessOutputSection { range })
} else {
None
}
})
.collect(),
}
}
pub fn deserialize(
saved_context: SavedContext,
path: Arc<Path>,
language_registry: Arc<LanguageRegistry>,
prompt_builder: Arc<PromptBuilder>,
slash_commands: Arc<SlashCommandWorkingSet>,
project: Option<Entity<Project>>,
telemetry: Option<Arc<Telemetry>>,
cx: &mut Context<Self>,
) -> Self {
let id = saved_context.id.clone().unwrap_or_else(ContextId::new);
let mut this = Self::new(
id,
ReplicaId::default(),
language::Capability::ReadWrite,
language_registry,
prompt_builder,
slash_commands,
project,
telemetry,
cx,
);
this.path = Some(path);
this.buffer.update(cx, |buffer, cx| {
buffer.set_text(saved_context.text.as_str(), cx)
});
let operations = saved_context.into_ops(&this.buffer, cx);
this.apply_ops(operations, cx);
this
}
pub fn id(&self) -> &ContextId {
&self.id
}
pub fn replica_id(&self) -> ReplicaId {
self.timestamp.replica_id
}
pub fn version(&self, cx: &App) -> ContextVersion {
ContextVersion {
context: self.version.clone(),
buffer: self.buffer.read(cx).version(),
}
}
pub fn slash_commands(&self) -> &Arc<SlashCommandWorkingSet> {
&self.slash_commands
}
pub fn set_capability(&mut self, capability: language::Capability, cx: &mut Context<Self>) {
self.buffer
.update(cx, |buffer, cx| buffer.set_capability(capability, cx));
}
fn next_timestamp(&mut self) -> clock::Lamport {
let timestamp = self.timestamp.tick();
self.version.observe(timestamp);
timestamp
}
pub fn serialize_ops(
&self,
since: &ContextVersion,
cx: &App,
) -> Task<Vec<proto::ContextOperation>> {
let buffer_ops = self
.buffer
.read(cx)
.serialize_ops(Some(since.buffer.clone()), cx);
let mut context_ops = self
.operations
.iter()
.filter(|op| !since.context.observed(op.timestamp()))
.cloned()
.collect::<Vec<_>>();
context_ops.extend(self.pending_ops.iter().cloned());
cx.background_spawn(async move {
let buffer_ops = buffer_ops.await;
context_ops.sort_unstable_by_key(|op| op.timestamp());
buffer_ops
.into_iter()
.map(|op| proto::ContextOperation {
variant: Some(proto::context_operation::Variant::BufferOperation(
proto::context_operation::BufferOperation {
operation: Some(op),
},
)),
})
.chain(context_ops.into_iter().map(|op| op.to_proto()))
.collect()
})
}
pub fn apply_ops(
&mut self,
ops: impl IntoIterator<Item = ContextOperation>,
cx: &mut Context<Self>,
) {
let mut buffer_ops = Vec::new();
for op in ops {
match op {
ContextOperation::BufferOperation(buffer_op) => buffer_ops.push(buffer_op),
op @ _ => self.pending_ops.push(op),
}
}
self.buffer
.update(cx, |buffer, cx| buffer.apply_ops(buffer_ops, cx));
self.flush_ops(cx);
}
fn flush_ops(&mut self, cx: &mut Context<AssistantContext>) {
let mut changed_messages = HashSet::default();
let mut summary_generated = false;
self.pending_ops.sort_unstable_by_key(|op| op.timestamp());
for op in mem::take(&mut self.pending_ops) {
if !self.can_apply_op(&op, cx) {
self.pending_ops.push(op);
continue;
}
let timestamp = op.timestamp();
match op.clone() {
ContextOperation::InsertMessage {
anchor, metadata, ..
} => {
if self.messages_metadata.contains_key(&anchor.id) {
// We already applied this operation.
} else {
changed_messages.insert(anchor.id);
self.insert_message(anchor, metadata, cx);
}
}
ContextOperation::UpdateMessage {
message_id,
metadata: new_metadata,
..
} => {
let metadata = self.messages_metadata.get_mut(&message_id).unwrap();
if new_metadata.timestamp > metadata.timestamp {
*metadata = new_metadata;
changed_messages.insert(message_id);
}
}
ContextOperation::UpdateSummary {
summary: new_summary,
..
} => {
if self.summary.timestamp().map_or(true, |current_timestamp| {
new_summary.timestamp > current_timestamp
}) {
self.summary = ContextSummary::Content(new_summary);
summary_generated = true;
}
}
ContextOperation::SlashCommandStarted {
id,
output_range,
name,
..
} => {
self.invoked_slash_commands.insert(
id,
InvokedSlashCommand {
name: name.into(),
range: output_range,
run_commands_in_ranges: Vec::new(),
status: InvokedSlashCommandStatus::Running(Task::ready(())),
transaction: None,
timestamp: id.0,
},
);
cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id: id });
}
ContextOperation::SlashCommandOutputSectionAdded { section, .. } => {
let buffer = self.buffer.read(cx);
if let Err(ix) = self
.slash_command_output_sections
.binary_search_by(|probe| probe.range.cmp(&section.range, buffer))
{
self.slash_command_output_sections
.insert(ix, section.clone());
cx.emit(ContextEvent::SlashCommandOutputSectionAdded { section });
}
}
ContextOperation::ThoughtProcessOutputSectionAdded { section, .. } => {
let buffer = self.buffer.read(cx);
if let Err(ix) = self
.thought_process_output_sections
.binary_search_by(|probe| probe.range.cmp(&section.range, buffer))
{
self.thought_process_output_sections
.insert(ix, section.clone());
}
}
ContextOperation::SlashCommandFinished {
id,
error_message,
timestamp,
..
} => {
if let Some(slash_command) = self.invoked_slash_commands.get_mut(&id) {
if timestamp > slash_command.timestamp {
slash_command.timestamp = timestamp;
match error_message {
Some(message) => {
slash_command.status =
InvokedSlashCommandStatus::Error(message.into());
}
None => {
slash_command.status = InvokedSlashCommandStatus::Finished;
}
}
cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id: id });
}
}
}
ContextOperation::BufferOperation(_) => unreachable!(),
}
self.version.observe(timestamp);
self.timestamp.observe(timestamp);
self.operations.push(op);
}
if !changed_messages.is_empty() {
self.message_roles_updated(changed_messages, cx);
cx.emit(ContextEvent::MessagesEdited);
cx.notify();
}
if summary_generated {
cx.emit(ContextEvent::SummaryChanged);
cx.emit(ContextEvent::SummaryGenerated);
cx.notify();
}
}
fn can_apply_op(&self, op: &ContextOperation, cx: &App) -> bool {
if !self.version.observed_all(op.version()) {
return false;
}
match op {
ContextOperation::InsertMessage { anchor, .. } => self
.buffer
.read(cx)
.version
.observed(anchor.start.timestamp),
ContextOperation::UpdateMessage { message_id, .. } => {
self.messages_metadata.contains_key(message_id)
}
ContextOperation::UpdateSummary { .. } => true,
ContextOperation::SlashCommandStarted { output_range, .. } => {
self.has_received_operations_for_anchor_range(output_range.clone(), cx)
}
ContextOperation::SlashCommandOutputSectionAdded { section, .. } => {
self.has_received_operations_for_anchor_range(section.range.clone(), cx)
}
ContextOperation::ThoughtProcessOutputSectionAdded { section, .. } => {
self.has_received_operations_for_anchor_range(section.range.clone(), cx)
}
ContextOperation::SlashCommandFinished { .. } => true,
ContextOperation::BufferOperation(_) => {
panic!("buffer operations should always be applied")
}
}
}
fn has_received_operations_for_anchor_range(
&self,
range: Range<text::Anchor>,
cx: &App,
) -> bool {
let version = &self.buffer.read(cx).version;
let observed_start = range.start == language::Anchor::MIN
|| range.start == language::Anchor::MAX
|| version.observed(range.start.timestamp);
let observed_end = range.end == language::Anchor::MIN
|| range.end == language::Anchor::MAX
|| version.observed(range.end.timestamp);
observed_start && observed_end
}
fn push_op(&mut self, op: ContextOperation, cx: &mut Context<Self>) {
self.operations.push(op.clone());
cx.emit(ContextEvent::Operation(op));
}
pub fn buffer(&self) -> &Entity<Buffer> {
&self.buffer
}
pub fn language_registry(&self) -> Arc<LanguageRegistry> {
self.language_registry.clone()
}
pub fn project(&self) -> Option<Entity<Project>> {
self.project.clone()
}
pub fn prompt_builder(&self) -> Arc<PromptBuilder> {
self.prompt_builder.clone()
}
pub fn path(&self) -> Option<&Arc<Path>> {
self.path.as_ref()
}
pub fn summary(&self) -> &ContextSummary {
&self.summary
}
pub fn parsed_slash_commands(&self) -> &[ParsedSlashCommand] {
&self.parsed_slash_commands
}
pub fn invoked_slash_command(
&self,
command_id: &InvokedSlashCommandId,
) -> Option<&InvokedSlashCommand> {
self.invoked_slash_commands.get(command_id)
}
pub fn slash_command_output_sections(&self) -> &[SlashCommandOutputSection<language::Anchor>] {
&self.slash_command_output_sections
}
pub fn thought_process_output_sections(
&self,
) -> &[ThoughtProcessOutputSection<language::Anchor>] {
&self.thought_process_output_sections
}
pub fn contains_files(&self, cx: &App) -> bool {
let buffer = self.buffer.read(cx);
self.slash_command_output_sections.iter().any(|section| {
section.is_valid(buffer)
&& section
.metadata
.as_ref()
.and_then(|metadata| {
serde_json::from_value::<FileCommandMetadata>(metadata.clone()).ok()
})
.is_some()
})
}
fn set_language(&mut self, cx: &mut Context<Self>) {
let markdown = self.language_registry.language_for_name("Markdown");
cx.spawn(async move |this, cx| {
let markdown = markdown.await?;
this.update(cx, |this, cx| {
this.buffer
.update(cx, |buffer, cx| buffer.set_language(Some(markdown), cx));
})
})
.detach_and_log_err(cx);
}
fn handle_buffer_event(
&mut self,
_: Entity<Buffer>,
event: &language::BufferEvent,
cx: &mut Context<Self>,
) {
match event {
language::BufferEvent::Operation {
operation,
is_local: true,
} => cx.emit(ContextEvent::Operation(ContextOperation::BufferOperation(
operation.clone(),
))),
language::BufferEvent::Edited => {
self.count_remaining_tokens(cx);
self.reparse(cx);
cx.emit(ContextEvent::MessagesEdited);
}
_ => {}
}
}
pub fn token_count(&self) -> Option<u64> {
self.token_count
}
pub(crate) fn count_remaining_tokens(&mut self, cx: &mut Context<Self>) {
// Assume it will be a Chat request, even though that takes fewer tokens (and risks going over the limit),
// because otherwise you see in the UI that your empty message has a bunch of tokens already used.
let Some(model) = LanguageModelRegistry::read_global(cx).default_model() else {
return;
};
let request = self.to_completion_request(Some(&model.model), cx);
let debounce = self.token_count.is_some();
self.pending_token_count = cx.spawn(async move |this, cx| {
async move {
if debounce {
cx.background_executor()
.timer(Duration::from_millis(200))
.await;
}
let token_count = cx
.update(|cx| model.model.count_tokens(request, cx))?
.await?;
this.update(cx, |this, cx| {
this.token_count = Some(token_count);
this.start_cache_warming(&model.model, cx);
cx.notify()
})
}
.log_err()
.await
});
}
pub fn mark_cache_anchors(
&mut self,
cache_configuration: &Option<LanguageModelCacheConfiguration>,
speculative: bool,
cx: &mut Context<Self>,
) -> bool {
let cache_configuration =
cache_configuration
.as_ref()
.unwrap_or(&LanguageModelCacheConfiguration {
max_cache_anchors: 0,
should_speculate: false,
min_total_token: 0,
});
let messages: Vec<Message> = self.messages(cx).collect();
let mut sorted_messages = messages.clone();
if speculative {
// Avoid caching the last message if this is a speculative cache fetch as
// it's likely to change.
sorted_messages.pop();
}
sorted_messages.retain(|m| m.role == Role::User);
sorted_messages.sort_by(|a, b| b.offset_range.len().cmp(&a.offset_range.len()));
let cache_anchors = if self.token_count.unwrap_or(0) < cache_configuration.min_total_token {
// If we have't hit the minimum threshold to enable caching, don't cache anything.
0
} else {
// Save 1 anchor for the inline assistant to use.
max(cache_configuration.max_cache_anchors, 1) - 1
};
sorted_messages.truncate(cache_anchors);
let anchors: HashSet<MessageId> = sorted_messages
.into_iter()
.map(|message| message.id)
.collect();
let buffer = self.buffer.read(cx).snapshot();
let invalidated_caches: HashSet<MessageId> = messages
.iter()
.scan(false, |encountered_invalid, message| {
let message_id = message.id;
let is_invalid = self
.messages_metadata
.get(&message_id)
.map_or(true, |metadata| {
!metadata.is_cache_valid(&buffer, &message.offset_range)
|| *encountered_invalid
});
*encountered_invalid |= is_invalid;
Some(if is_invalid { Some(message_id) } else { None })
})
.flatten()
.collect();
let last_anchor = messages.iter().rev().find_map(|message| {
if anchors.contains(&message.id) {
Some(message.id)
} else {
None
}
});
let mut new_anchor_needs_caching = false;
let current_version = &buffer.version;
// If we have no anchors, mark all messages as not being cached.
let mut hit_last_anchor = last_anchor.is_none();
for message in messages.iter() {
if hit_last_anchor {
self.update_metadata(message.id, cx, |metadata| metadata.cache = None);
continue;
}
if let Some(last_anchor) = last_anchor {
if message.id == last_anchor {
hit_last_anchor = true;
}
}
new_anchor_needs_caching = new_anchor_needs_caching
|| (invalidated_caches.contains(&message.id) && anchors.contains(&message.id));
self.update_metadata(message.id, cx, |metadata| {
let cache_status = if invalidated_caches.contains(&message.id) {
CacheStatus::Pending
} else {
metadata
.cache
.as_ref()
.map_or(CacheStatus::Pending, |cm| cm.status.clone())
};
metadata.cache = Some(MessageCacheMetadata {
is_anchor: anchors.contains(&message.id),
is_final_anchor: hit_last_anchor,
status: cache_status,
cached_at: current_version.clone(),
});
});
}
new_anchor_needs_caching
}
fn start_cache_warming(&mut self, model: &Arc<dyn LanguageModel>, cx: &mut Context<Self>) {
let cache_configuration = model.cache_configuration();
if !self.mark_cache_anchors(&cache_configuration, true, cx) {
return;
}
if !self.pending_completions.is_empty() {
return;
}
if let Some(cache_configuration) = cache_configuration {
if !cache_configuration.should_speculate {
return;
}
}
let request = {
let mut req = self.to_completion_request(Some(&model), cx);
// Skip the last message because it's likely to change and
// therefore would be a waste to cache.
req.messages.pop();
req.messages.push(LanguageModelRequestMessage {
role: Role::User,
content: vec!["Respond only with OK, nothing else.".into()],
cache: false,
});
req
};
let model = Arc::clone(model);
self.pending_cache_warming_task = cx.spawn(async move |this, cx| {
async move {
match model.stream_completion(request, &cx).await {
Ok(mut stream) => {
stream.next().await;
log::info!("Cache warming completed successfully");
}
Err(e) => {
log::warn!("Cache warming failed: {}", e);
}
};
this.update(cx, |this, cx| {
this.update_cache_status_for_completion(cx);
})
.ok();
anyhow::Ok(())
}
.log_err()
.await
});
}
pub fn update_cache_status_for_completion(&mut self, cx: &mut Context<Self>) {
let cached_message_ids: Vec<MessageId> = self
.messages_metadata
.iter()
.filter_map(|(message_id, metadata)| {
metadata.cache.as_ref().and_then(|cache| {
if cache.status == CacheStatus::Pending {
Some(*message_id)
} else {
None
}
})
})
.collect();
for message_id in cached_message_ids {
self.update_metadata(message_id, cx, |metadata| {
if let Some(cache) = &mut metadata.cache {
cache.status = CacheStatus::Cached;
}
});
}
cx.notify();
}
pub fn reparse(&mut self, cx: &mut Context<Self>) {
let buffer = self.buffer.read(cx).text_snapshot();
let mut row_ranges = self
.edits_since_last_parse
.consume()
.into_iter()
.map(|edit| {
let start_row = buffer.offset_to_point(edit.new.start).row;
let end_row = buffer.offset_to_point(edit.new.end).row + 1;
start_row..end_row
})
.peekable();
let mut removed_parsed_slash_command_ranges = Vec::new();
let mut updated_parsed_slash_commands = Vec::new();
while let Some(mut row_range) = row_ranges.next() {
while let Some(next_row_range) = row_ranges.peek() {
if row_range.end >= next_row_range.start {
row_range.end = next_row_range.end;
row_ranges.next();
} else {
break;
}
}
let start = buffer.anchor_before(Point::new(row_range.start, 0));
let end = buffer.anchor_after(Point::new(
row_range.end - 1,
buffer.line_len(row_range.end - 1),
));
self.reparse_slash_commands_in_range(
start..end,
&buffer,
&mut updated_parsed_slash_commands,
&mut removed_parsed_slash_command_ranges,
cx,
);
self.invalidate_pending_slash_commands(&buffer, cx);
}
if !updated_parsed_slash_commands.is_empty()
|| !removed_parsed_slash_command_ranges.is_empty()
{
cx.emit(ContextEvent::ParsedSlashCommandsUpdated {
removed: removed_parsed_slash_command_ranges,
updated: updated_parsed_slash_commands,
});
}
}
fn reparse_slash_commands_in_range(
&mut self,
range: Range<text::Anchor>,
buffer: &BufferSnapshot,
updated: &mut Vec<ParsedSlashCommand>,
removed: &mut Vec<Range<text::Anchor>>,
cx: &App,
) {
let old_range = self.pending_command_indices_for_range(range.clone(), cx);
let mut new_commands = Vec::new();
let mut lines = buffer.text_for_range(range).lines();
let mut offset = lines.offset();
while let Some(line) = lines.next() {
if let Some(command_line) = SlashCommandLine::parse(line) {
let name = &line[command_line.name.clone()];
let arguments = command_line
.arguments
.iter()
.filter_map(|argument_range| {
if argument_range.is_empty() {
None
} else {
line.get(argument_range.clone())
}
})
.map(ToOwned::to_owned)
.collect::<SmallVec<_>>();
if let Some(command) = self.slash_commands.command(name, cx) {
if !command.requires_argument() || !arguments.is_empty() {
let start_ix = offset + command_line.name.start - 1;
let end_ix = offset
+ command_line
.arguments
.last()
.map_or(command_line.name.end, |argument| argument.end);
let source_range =
buffer.anchor_after(start_ix)..buffer.anchor_after(end_ix);
let pending_command = ParsedSlashCommand {
name: name.to_string(),
arguments,
source_range,
status: PendingSlashCommandStatus::Idle,
};
updated.push(pending_command.clone());
new_commands.push(pending_command);
}
}
}
offset = lines.offset();
}
let removed_commands = self.parsed_slash_commands.splice(old_range, new_commands);
removed.extend(removed_commands.map(|command| command.source_range));
}
fn invalidate_pending_slash_commands(
&mut self,
buffer: &BufferSnapshot,
cx: &mut Context<Self>,
) {
let mut invalidated_command_ids = Vec::new();
for (&command_id, command) in self.invoked_slash_commands.iter_mut() {
if !matches!(command.status, InvokedSlashCommandStatus::Finished)
&& (!command.range.start.is_valid(buffer) || !command.range.end.is_valid(buffer))
{
command.status = InvokedSlashCommandStatus::Finished;
cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id });
invalidated_command_ids.push(command_id);
}
}
for command_id in invalidated_command_ids {
let version = self.version.clone();
let timestamp = self.next_timestamp();
self.push_op(
ContextOperation::SlashCommandFinished {
id: command_id,
timestamp,
error_message: None,
version: version.clone(),
},
cx,
);
}
}
pub fn pending_command_for_position(
&mut self,
position: language::Anchor,
cx: &mut Context<Self>,
) -> Option<&mut ParsedSlashCommand> {
let buffer = self.buffer.read(cx);
match self
.parsed_slash_commands
.binary_search_by(|probe| probe.source_range.end.cmp(&position, buffer))
{
Ok(ix) => Some(&mut self.parsed_slash_commands[ix]),
Err(ix) => {
let cmd = self.parsed_slash_commands.get_mut(ix)?;
if position.cmp(&cmd.source_range.start, buffer).is_ge()
&& position.cmp(&cmd.source_range.end, buffer).is_le()
{
Some(cmd)
} else {
None
}
}
}
}
pub fn pending_commands_for_range(
&self,
range: Range<language::Anchor>,
cx: &App,
) -> &[ParsedSlashCommand] {
let range = self.pending_command_indices_for_range(range, cx);
&self.parsed_slash_commands[range]
}
fn pending_command_indices_for_range(
&self,
range: Range<language::Anchor>,
cx: &App,
) -> Range<usize> {
self.indices_intersecting_buffer_range(&self.parsed_slash_commands, range, cx)
}
fn indices_intersecting_buffer_range<T: ContextAnnotation>(
&self,
all_annotations: &[T],
range: Range<language::Anchor>,
cx: &App,
) -> Range<usize> {
let buffer = self.buffer.read(cx);
let start_ix = match all_annotations
.binary_search_by(|probe| probe.range().end.cmp(&range.start, &buffer))
{
Ok(ix) | Err(ix) => ix,
};
let end_ix = match all_annotations
.binary_search_by(|probe| probe.range().start.cmp(&range.end, &buffer))
{
Ok(ix) => ix + 1,
Err(ix) => ix,
};
start_ix..end_ix
}
pub fn insert_command_output(
&mut self,
command_source_range: Range<language::Anchor>,
name: &str,
output: Task<SlashCommandResult>,
ensure_trailing_newline: bool,
cx: &mut Context<Self>,
) {
let version = self.version.clone();
let command_id = InvokedSlashCommandId(self.next_timestamp());
const PENDING_OUTPUT_END_MARKER: &str = "…";
let (command_range, command_source_range, insert_position, first_transaction) =
self.buffer.update(cx, |buffer, cx| {
let command_source_range = command_source_range.to_offset(buffer);
let mut insertion = format!("\n{PENDING_OUTPUT_END_MARKER}");
if ensure_trailing_newline {
insertion.push('\n');
}
buffer.finalize_last_transaction();
buffer.start_transaction();
buffer.edit(
[(
command_source_range.end..command_source_range.end,
insertion,
)],
None,
cx,
);
let first_transaction = buffer.end_transaction(cx).unwrap();
buffer.finalize_last_transaction();
let insert_position = buffer.anchor_after(command_source_range.end + 1);
let command_range = buffer.anchor_after(command_source_range.start)
..buffer.anchor_before(
command_source_range.end + 1 + PENDING_OUTPUT_END_MARKER.len(),
);
let command_source_range = buffer.anchor_before(command_source_range.start)
..buffer.anchor_before(command_source_range.end + 1);
(
command_range,
command_source_range,
insert_position,
first_transaction,
)
});
self.reparse(cx);
let insert_output_task = cx.spawn(async move |this, cx| {
let run_command = async {
let mut stream = output.await?;
struct PendingSection {
start: language::Anchor,
icon: IconName,
label: SharedString,
metadata: Option<serde_json::Value>,
}
let mut pending_section_stack: Vec<PendingSection> = Vec::new();
let mut last_role: Option<Role> = None;
let mut last_section_range = None;
while let Some(event) = stream.next().await {
let event = event?;
this.update(cx, |this, cx| {
this.buffer.update(cx, |buffer, _cx| {
buffer.finalize_last_transaction();
buffer.start_transaction()
});
match event {
SlashCommandEvent::StartMessage {
role,
merge_same_roles,
} => {
if !merge_same_roles && Some(role) != last_role {
let buffer = this.buffer.read(cx);
let offset = insert_position.to_offset(buffer);
this.insert_message_at_offset(
offset,
role,
MessageStatus::Pending,
cx,
);
}
last_role = Some(role);
}
SlashCommandEvent::StartSection {
icon,
label,
metadata,
} => {
this.buffer.update(cx, |buffer, cx| {
let insert_point = insert_position.to_point(buffer);
if insert_point.column > 0 {
buffer.edit([(insert_point..insert_point, "\n")], None, cx);
}
pending_section_stack.push(PendingSection {
start: buffer.anchor_before(insert_position),
icon,
label,
metadata,
});
});
}
SlashCommandEvent::Content(SlashCommandContent::Text {
text,
run_commands_in_text,
}) => {
let start = this.buffer.read(cx).anchor_before(insert_position);
this.buffer.update(cx, |buffer, cx| {
buffer.edit(
[(insert_position..insert_position, text)],
None,
cx,
)
});
let end = this.buffer.read(cx).anchor_before(insert_position);
if run_commands_in_text {
if let Some(invoked_slash_command) =
this.invoked_slash_commands.get_mut(&command_id)
{
invoked_slash_command
.run_commands_in_ranges
.push(start..end);
}
}
}
SlashCommandEvent::EndSection => {
if let Some(pending_section) = pending_section_stack.pop() {
let offset_range = (pending_section.start..insert_position)
.to_offset(this.buffer.read(cx));
if !offset_range.is_empty() {
let range = this.buffer.update(cx, |buffer, _cx| {
buffer.anchor_after(offset_range.start)
..buffer.anchor_before(offset_range.end)
});
this.insert_slash_command_output_section(
SlashCommandOutputSection {
range: range.clone(),
icon: pending_section.icon,
label: pending_section.label,
metadata: pending_section.metadata,
},
cx,
);
last_section_range = Some(range);
}
}
}
}
this.buffer.update(cx, |buffer, cx| {
if let Some(event_transaction) = buffer.end_transaction(cx) {
buffer.merge_transactions(event_transaction, first_transaction);
}
});
})?;
}
this.update(cx, |this, cx| {
this.buffer.update(cx, |buffer, cx| {
buffer.finalize_last_transaction();
buffer.start_transaction();
let mut deletions = vec![(command_source_range.to_offset(buffer), "")];
let insert_position = insert_position.to_offset(buffer);
let command_range_end = command_range.end.to_offset(buffer);
if buffer.contains_str_at(insert_position, PENDING_OUTPUT_END_MARKER) {
deletions.push((
insert_position..insert_position + PENDING_OUTPUT_END_MARKER.len(),
"",
));
}
if ensure_trailing_newline
&& buffer.contains_str_at(command_range_end, "\n")
{
let newline_offset = insert_position.saturating_sub(1);
if buffer.contains_str_at(newline_offset, "\n")
&& last_section_range.map_or(true, |last_section_range| {
!last_section_range
.to_offset(buffer)
.contains(&newline_offset)
})
{
deletions.push((command_range_end..command_range_end + 1, ""));
}
}
buffer.edit(deletions, None, cx);
if let Some(deletion_transaction) = buffer.end_transaction(cx) {
buffer.merge_transactions(deletion_transaction, first_transaction);
}
});
})?;
debug_assert!(pending_section_stack.is_empty());
anyhow::Ok(())
};
let command_result = run_command.await;
this.update(cx, |this, cx| {
let version = this.version.clone();
let timestamp = this.next_timestamp();
let Some(invoked_slash_command) = this.invoked_slash_commands.get_mut(&command_id)
else {
return;
};
let mut error_message = None;
match command_result {
Ok(()) => {
invoked_slash_command.status = InvokedSlashCommandStatus::Finished;
}
Err(error) => {
let message = error.to_string();
invoked_slash_command.status =
InvokedSlashCommandStatus::Error(message.clone().into());
error_message = Some(message);
}
}
cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id });
this.push_op(
ContextOperation::SlashCommandFinished {
id: command_id,
timestamp,
error_message,
version,
},
cx,
);
})
.ok();
});
self.invoked_slash_commands.insert(
command_id,
InvokedSlashCommand {
name: name.to_string().into(),
range: command_range.clone(),
run_commands_in_ranges: Vec::new(),
status: InvokedSlashCommandStatus::Running(insert_output_task),
transaction: Some(first_transaction),
timestamp: command_id.0,
},
);
cx.emit(ContextEvent::InvokedSlashCommandChanged { command_id });
self.push_op(
ContextOperation::SlashCommandStarted {
id: command_id,
output_range: command_range,
name: name.to_string(),
version,
},
cx,
);
}
fn insert_slash_command_output_section(
&mut self,
section: SlashCommandOutputSection<language::Anchor>,
cx: &mut Context<Self>,
) {
let buffer = self.buffer.read(cx);
let insertion_ix = match self
.slash_command_output_sections
.binary_search_by(|probe| probe.range.cmp(&section.range, buffer))
{
Ok(ix) | Err(ix) => ix,
};
self.slash_command_output_sections
.insert(insertion_ix, section.clone());
cx.emit(ContextEvent::SlashCommandOutputSectionAdded {
section: section.clone(),
});
let version = self.version.clone();
let timestamp = self.next_timestamp();
self.push_op(
ContextOperation::SlashCommandOutputSectionAdded {
timestamp,
section,
version,
},
cx,
);
}
fn insert_thought_process_output_section(
&mut self,
section: ThoughtProcessOutputSection<language::Anchor>,
cx: &mut Context<Self>,
) {
let buffer = self.buffer.read(cx);
let insertion_ix = match self
.thought_process_output_sections
.binary_search_by(|probe| probe.range.cmp(&section.range, buffer))
{
Ok(ix) | Err(ix) => ix,
};
self.thought_process_output_sections
.insert(insertion_ix, section.clone());
// cx.emit(ContextEvent::ThoughtProcessOutputSectionAdded {
// section: section.clone(),
// });
let version = self.version.clone();
let timestamp = self.next_timestamp();
self.push_op(
ContextOperation::ThoughtProcessOutputSectionAdded {
timestamp,
section,
version,
},
cx,
);
}
pub fn completion_provider_changed(&mut self, cx: &mut Context<Self>) {
self.count_remaining_tokens(cx);
}
fn get_last_valid_message_id(&self, cx: &Context<Self>) -> Option<MessageId> {
self.message_anchors.iter().rev().find_map(|message| {
message
.start
.is_valid(self.buffer.read(cx))
.then_some(message.id)
})
}
pub fn assist(&mut self, cx: &mut Context<Self>) -> Option<MessageAnchor> {
let model_registry = LanguageModelRegistry::read_global(cx);
let model = model_registry.default_model()?;
let last_message_id = self.get_last_valid_message_id(cx)?;
if !model.provider.is_authenticated(cx) {
log::info!("completion provider has no credentials");
return None;
}
let model = model.model;
// Compute which messages to cache, including the last one.
self.mark_cache_anchors(&model.cache_configuration(), false, cx);
let request = self.to_completion_request(Some(&model), cx);
let assistant_message = self
.insert_message_after(last_message_id, Role::Assistant, MessageStatus::Pending, cx)
.unwrap();
// Queue up the user's next reply.
let user_message = self
.insert_message_after(assistant_message.id, Role::User, MessageStatus::Done, cx)
.unwrap();
let pending_completion_id = post_inc(&mut self.completion_count);
let task = cx.spawn({
async move |this, cx| {
let stream = model.stream_completion(request, &cx);
let assistant_message_id = assistant_message.id;
let mut response_latency = None;
let stream_completion = async {
let request_start = Instant::now();
let mut events = stream.await?;
let mut stop_reason = StopReason::EndTurn;
let mut thought_process_stack = Vec::new();
const THOUGHT_PROCESS_START_MARKER: &str = "<think>\n";
const THOUGHT_PROCESS_END_MARKER: &str = "\n</think>";
while let Some(event) = events.next().await {
if response_latency.is_none() {
response_latency = Some(request_start.elapsed());
}
let event = event?;
let mut context_event = None;
let mut thought_process_output_section = None;
this.update(cx, |this, cx| {
let message_ix = this
.message_anchors
.iter()
.position(|message| message.id == assistant_message_id)?;
this.buffer.update(cx, |buffer, cx| {
let message_old_end_offset = this.message_anchors[message_ix + 1..]
.iter()
.find(|message| message.start.is_valid(buffer))
.map_or(buffer.len(), |message| {
message.start.to_offset(buffer).saturating_sub(1)
});
match event {
LanguageModelCompletionEvent::StatusUpdate { .. } => {}
LanguageModelCompletionEvent::StartMessage { .. } => {}
LanguageModelCompletionEvent::Stop(reason) => {
stop_reason = reason;
}
LanguageModelCompletionEvent::Thinking { text: chunk, .. } => {
if thought_process_stack.is_empty() {
let start =
buffer.anchor_before(message_old_end_offset);
thought_process_stack.push(start);
let chunk =
format!("{THOUGHT_PROCESS_START_MARKER}{chunk}{THOUGHT_PROCESS_END_MARKER}");
let chunk_len = chunk.len();
buffer.edit(
[(
message_old_end_offset..message_old_end_offset,
chunk,
)],
None,
cx,
);
let end = buffer
.anchor_before(message_old_end_offset + chunk_len);
context_event = Some(
ContextEvent::StartedThoughtProcess(start..end),
);
} else {
// This ensures that all the thinking chunks are inserted inside the thinking tag
let insertion_position =
message_old_end_offset - THOUGHT_PROCESS_END_MARKER.len();
buffer.edit(
[(insertion_position..insertion_position, chunk)],
None,
cx,
);
}
}
LanguageModelCompletionEvent::RedactedThinking { .. } => {},
LanguageModelCompletionEvent::Text(mut chunk) => {
if let Some(start) = thought_process_stack.pop() {
let end = buffer.anchor_before(message_old_end_offset);
context_event =
Some(ContextEvent::EndedThoughtProcess(end));
thought_process_output_section =
Some(ThoughtProcessOutputSection {
range: start..end,
});
chunk.insert_str(0, "\n\n");
}
buffer.edit(
[(
message_old_end_offset..message_old_end_offset,
chunk,
)],
None,
cx,
);
}
LanguageModelCompletionEvent::ToolUse(_) |
LanguageModelCompletionEvent::ToolUseJsonParseError { .. } |
LanguageModelCompletionEvent::UsageUpdate(_) => {}
}
});
if let Some(section) = thought_process_output_section.take() {
this.insert_thought_process_output_section(section, cx);
}
if let Some(context_event) = context_event.take() {
cx.emit(context_event);
}
cx.emit(ContextEvent::StreamedCompletion);
Some(())
})?;
smol::future::yield_now().await;
}
this.update(cx, |this, cx| {
this.pending_completions
.retain(|completion| completion.id != pending_completion_id);
this.summarize(false, cx);
this.update_cache_status_for_completion(cx);
})?;
anyhow::Ok(stop_reason)
};
let result = stream_completion.await;
this.update(cx, |this, cx| {
let error_message = if let Some(error) = result.as_ref().err() {
if error.is::<PaymentRequiredError>() {
cx.emit(ContextEvent::ShowPaymentRequiredError);
this.update_metadata(assistant_message_id, cx, |metadata| {
metadata.status = MessageStatus::Canceled;
});
Some(error.to_string())
} else {
let error_message = error
.chain()
.map(|err| err.to_string())
.collect::<Vec<_>>()
.join("\n");
cx.emit(ContextEvent::ShowAssistError(SharedString::from(
error_message.clone(),
)));
this.update_metadata(assistant_message_id, cx, |metadata| {
metadata.status =
MessageStatus::Error(SharedString::from(error_message.clone()));
});
Some(error_message)
}
} else {
this.update_metadata(assistant_message_id, cx, |metadata| {
metadata.status = MessageStatus::Done;
});
None
};
let language_name = this
.buffer
.read(cx)
.language()
.map(|language| language.name());
report_assistant_event(
AssistantEventData {
conversation_id: Some(this.id.0.clone()),
kind: AssistantKind::Panel,
phase: AssistantPhase::Response,
message_id: None,
model: model.telemetry_id(),
model_provider: model.provider_id().to_string(),
response_latency,
error_message,
language_name: language_name.map(|name| name.to_proto()),
},
this.telemetry.clone(),
cx.http_client(),
model.api_key(cx),
cx.background_executor(),
);
if let Ok(stop_reason) = result {
match stop_reason {
StopReason::ToolUse => {}
StopReason::EndTurn => {}
StopReason::MaxTokens => {}
StopReason::Refusal => {}
}
}
})
.ok();
}
});
self.pending_completions.push(PendingCompletion {
id: pending_completion_id,
assistant_message_id: assistant_message.id,
_task: task,
});
Some(user_message)
}
pub fn to_xml(&self, cx: &App) -> String {
let mut output = String::new();
let buffer = self.buffer.read(cx);
for message in self.messages(cx) {
if message.status != MessageStatus::Done {
continue;
}
writeln!(&mut output, "<{}>", message.role).unwrap();
for chunk in buffer.text_for_range(message.offset_range) {
output.push_str(chunk);
}
if !output.ends_with('\n') {
output.push('\n');
}
writeln!(&mut output, "</{}>", message.role).unwrap();
}
output
}
pub fn to_completion_request(
&self,
model: Option<&Arc<dyn LanguageModel>>,
cx: &App,
) -> LanguageModelRequest {
let buffer = self.buffer.read(cx);
let mut contents = self.contents(cx).peekable();
fn collect_text_content(buffer: &Buffer, range: Range<usize>) -> Option<String> {
let text: String = buffer.text_for_range(range.clone()).collect();
if text.trim().is_empty() {
None
} else {
Some(text)
}
}
let mut completion_request = LanguageModelRequest {
thread_id: None,
prompt_id: None,
intent: Some(CompletionIntent::UserPrompt),
mode: None,
messages: Vec::new(),
tools: Vec::new(),
tool_choice: None,
stop: Vec::new(),
temperature: model.and_then(|model| AgentSettings::temperature_for_model(model, cx)),
thinking_allowed: true,
};
for message in self.messages(cx) {
if message.status != MessageStatus::Done {
continue;
}
let mut offset = message.offset_range.start;
let mut request_message = LanguageModelRequestMessage {
role: message.role,
content: Vec::new(),
cache: message
.cache
.as_ref()
.map_or(false, |cache| cache.is_anchor),
};
while let Some(content) = contents.peek() {
if content
.range()
.end
.cmp(&message.anchor_range.end, buffer)
.is_lt()
{
let content = contents.next().unwrap();
let range = content.range().to_offset(buffer);
request_message.content.extend(
collect_text_content(buffer, offset..range.start).map(MessageContent::Text),
);
match content {
Content::Image { image, .. } => {
if let Some(image) = image.clone().now_or_never().flatten() {
request_message
.content
.push(language_model::MessageContent::Image(image));
}
}
}
offset = range.end;
} else {
break;
}
}
request_message.content.extend(
collect_text_content(buffer, offset..message.offset_range.end)
.map(MessageContent::Text),
);
if !request_message.contents_empty() {
completion_request.messages.push(request_message);
}
}
let supports_burn_mode = if let Some(model) = model {
model.supports_burn_mode()
} else {
false
};
if supports_burn_mode {
completion_request.mode = Some(self.completion_mode.into());
}
completion_request
}
pub fn cancel_last_assist(&mut self, cx: &mut Context<Self>) -> bool {
if let Some(pending_completion) = self.pending_completions.pop() {
self.update_metadata(pending_completion.assistant_message_id, cx, |metadata| {
if metadata.status == MessageStatus::Pending {
metadata.status = MessageStatus::Canceled;
}
});
true
} else {
false
}
}
pub fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut Context<Self>) {
for id in &ids {
if let Some(metadata) = self.messages_metadata.get(id) {
let role = metadata.role.cycle();
self.update_metadata(*id, cx, |metadata| metadata.role = role);
}
}
self.message_roles_updated(ids, cx);
}
fn message_roles_updated(&mut self, ids: HashSet<MessageId>, cx: &mut Context<Self>) {
let mut ranges = Vec::new();
for message in self.messages(cx) {
if ids.contains(&message.id) {
ranges.push(message.anchor_range.clone());
}
}
}
pub fn update_metadata(
&mut self,
id: MessageId,
cx: &mut Context<Self>,
f: impl FnOnce(&mut MessageMetadata),
) {
let version = self.version.clone();
let timestamp = self.next_timestamp();
if let Some(metadata) = self.messages_metadata.get_mut(&id) {
f(metadata);
metadata.timestamp = timestamp;
let operation = ContextOperation::UpdateMessage {
message_id: id,
metadata: metadata.clone(),
version,
};
self.push_op(operation, cx);
cx.emit(ContextEvent::MessagesEdited);
cx.notify();
}
}
pub fn insert_message_after(
&mut self,
message_id: MessageId,
role: Role,
status: MessageStatus,
cx: &mut Context<Self>,
) -> Option<MessageAnchor> {
if let Some(prev_message_ix) = self
.message_anchors
.iter()
.position(|message| message.id == message_id)
{
// Find the next valid message after the one we were given.
let mut next_message_ix = prev_message_ix + 1;
while let Some(next_message) = self.message_anchors.get(next_message_ix) {
if next_message.start.is_valid(self.buffer.read(cx)) {
break;
}
next_message_ix += 1;
}
let buffer = self.buffer.read(cx);
let offset = self
.message_anchors
.get(next_message_ix)
.map_or(buffer.len(), |message| {
buffer.clip_offset(message.start.to_offset(buffer) - 1, Bias::Left)
});
Some(self.insert_message_at_offset(offset, role, status, cx))
} else {
None
}
}
fn insert_message_at_offset(
&mut self,
offset: usize,
role: Role,
status: MessageStatus,
cx: &mut Context<Self>,
) -> MessageAnchor {
let start = self.buffer.update(cx, |buffer, cx| {
buffer.edit([(offset..offset, "\n")], None, cx);
buffer.anchor_before(offset + 1)
});
let version = self.version.clone();
let anchor = MessageAnchor {
id: MessageId(self.next_timestamp()),
start,
};
let metadata = MessageMetadata {
role,
status,
timestamp: anchor.id.0,
cache: None,
};
self.insert_message(anchor.clone(), metadata.clone(), cx);
self.push_op(
ContextOperation::InsertMessage {
anchor: anchor.clone(),
metadata,
version,
},
cx,
);
anchor
}
pub fn insert_content(&mut self, content: Content, cx: &mut Context<Self>) {
let buffer = self.buffer.read(cx);
let insertion_ix = match self
.contents
.binary_search_by(|probe| probe.cmp(&content, buffer))
{
Ok(ix) => {
self.contents.remove(ix);
ix
}
Err(ix) => ix,
};
self.contents.insert(insertion_ix, content);
cx.emit(ContextEvent::MessagesEdited);
}
pub fn contents<'a>(&'a self, cx: &'a App) -> impl 'a + Iterator<Item = Content> {
let buffer = self.buffer.read(cx);
self.contents
.iter()
.filter(|content| {
let range = content.range();
range.start.is_valid(buffer) && range.end.is_valid(buffer)
})
.cloned()
}
pub fn split_message(
&mut self,
range: Range<usize>,
cx: &mut Context<Self>,
) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
let start_message = self.message_for_offset(range.start, cx);
let end_message = self.message_for_offset(range.end, cx);
if let Some((start_message, end_message)) = start_message.zip(end_message) {
// Prevent splitting when range spans multiple messages.
if start_message.id != end_message.id {
return (None, None);
}
let message = start_message;
let at_end = range.end >= message.offset_range.end.saturating_sub(1);
let role_after = if range.start == range.end || at_end {
Role::User
} else {
message.role
};
let role = message.role;
let mut edited_buffer = false;
let mut suffix_start = None;
// TODO: why did this start panicking?
if range.start > message.offset_range.start
&& range.end < message.offset_range.end.saturating_sub(1)
{
if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
suffix_start = Some(range.end + 1);
} else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
suffix_start = Some(range.end);
}
}
let version = self.version.clone();
let suffix = if let Some(suffix_start) = suffix_start {
MessageAnchor {
id: MessageId(self.next_timestamp()),
start: self.buffer.read(cx).anchor_before(suffix_start),
}
} else {
self.buffer.update(cx, |buffer, cx| {
buffer.edit([(range.end..range.end, "\n")], None, cx);
});
edited_buffer = true;
MessageAnchor {
id: MessageId(self.next_timestamp()),
start: self.buffer.read(cx).anchor_before(range.end + 1),
}
};
let suffix_metadata = MessageMetadata {
role: role_after,
status: MessageStatus::Done,
timestamp: suffix.id.0,
cache: None,
};
self.insert_message(suffix.clone(), suffix_metadata.clone(), cx);
self.push_op(
ContextOperation::InsertMessage {
anchor: suffix.clone(),
metadata: suffix_metadata,
version,
},
cx,
);
let new_messages =
if range.start == range.end || range.start == message.offset_range.start {
(None, Some(suffix))
} else {
let mut prefix_end = None;
if range.start > message.offset_range.start
&& range.end < message.offset_range.end - 1
{
if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
prefix_end = Some(range.start + 1);
} else if self.buffer.read(cx).reversed_chars_at(range.start).next()
== Some('\n')
{
prefix_end = Some(range.start);
}
}
let version = self.version.clone();
let selection = if let Some(prefix_end) = prefix_end {
MessageAnchor {
id: MessageId(self.next_timestamp()),
start: self.buffer.read(cx).anchor_before(prefix_end),
}
} else {
self.buffer.update(cx, |buffer, cx| {
buffer.edit([(range.start..range.start, "\n")], None, cx)
});
edited_buffer = true;
MessageAnchor {
id: MessageId(self.next_timestamp()),
start: self.buffer.read(cx).anchor_before(range.end + 1),
}
};
let selection_metadata = MessageMetadata {
role,
status: MessageStatus::Done,
timestamp: selection.id.0,
cache: None,
};
self.insert_message(selection.clone(), selection_metadata.clone(), cx);
self.push_op(
ContextOperation::InsertMessage {
anchor: selection.clone(),
metadata: selection_metadata,
version,
},
cx,
);
(Some(selection), Some(suffix))
};
if !edited_buffer {
cx.emit(ContextEvent::MessagesEdited);
}
new_messages
} else {
(None, None)
}
}
fn insert_message(
&mut self,
new_anchor: MessageAnchor,
new_metadata: MessageMetadata,
cx: &mut Context<Self>,
) {
cx.emit(ContextEvent::MessagesEdited);
self.messages_metadata.insert(new_anchor.id, new_metadata);
let buffer = self.buffer.read(cx);
let insertion_ix = self
.message_anchors
.iter()
.position(|anchor| {
let comparison = new_anchor.start.cmp(&anchor.start, buffer);
comparison.is_lt() || (comparison.is_eq() && new_anchor.id > anchor.id)
})
.unwrap_or(self.message_anchors.len());
self.message_anchors.insert(insertion_ix, new_anchor);
}
pub fn summarize(&mut self, mut replace_old: bool, cx: &mut Context<Self>) {
let Some(model) = LanguageModelRegistry::read_global(cx).default_model() else {
return;
};
if replace_old || (self.message_anchors.len() >= 2 && self.summary.is_pending()) {
if !model.provider.is_authenticated(cx) {
return;
}
let mut request = self.to_completion_request(Some(&model.model), cx);
request.messages.push(LanguageModelRequestMessage {
role: Role::User,
content: vec![
"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:`"
.into(),
],
cache: false,
});
// If there is no summary, it is set with `done: false` so that "Loading Summary…" can
// be displayed.
match self.summary {
ContextSummary::Pending | ContextSummary::Error => {
self.summary = ContextSummary::Content(ContextSummaryContent {
text: "".to_string(),
done: false,
timestamp: clock::Lamport::default(),
});
replace_old = true;
}
ContextSummary::Content(_) => {}
}
self.summary_task = cx.spawn(async move |this, cx| {
let result = async {
let stream = model.model.stream_completion_text(request, &cx);
let mut messages = stream.await?;
let mut replaced = !replace_old;
while let Some(message) = messages.stream.next().await {
let text = message?;
let mut lines = text.lines();
this.update(cx, |this, cx| {
let version = this.version.clone();
let timestamp = this.next_timestamp();
let summary = this.summary.content_or_set_empty();
if !replaced && replace_old {
summary.text.clear();
replaced = true;
}
summary.text.extend(lines.next());
summary.timestamp = timestamp;
let operation = ContextOperation::UpdateSummary {
summary: summary.clone(),
version,
};
this.push_op(operation, cx);
cx.emit(ContextEvent::SummaryChanged);
cx.emit(ContextEvent::SummaryGenerated);
})?;
// Stop if the LLM generated multiple lines.
if lines.next().is_some() {
break;
}
}
this.read_with(cx, |this, _cx| {
if let Some(summary) = this.summary.content() {
if summary.text.is_empty() {
bail!("Model generated an empty summary");
}
}
Ok(())
})??;
this.update(cx, |this, cx| {
let version = this.version.clone();
let timestamp = this.next_timestamp();
if let Some(summary) = this.summary.content_as_mut() {
summary.done = true;
summary.timestamp = timestamp;
let operation = ContextOperation::UpdateSummary {
summary: summary.clone(),
version,
};
this.push_op(operation, cx);
cx.emit(ContextEvent::SummaryChanged);
cx.emit(ContextEvent::SummaryGenerated);
}
})?;
anyhow::Ok(())
}
.await;
if let Err(err) = result {
this.update(cx, |this, cx| {
this.summary = ContextSummary::Error;
cx.emit(ContextEvent::SummaryChanged);
})
.log_err();
log::error!("Error generating context summary: {}", err);
}
Some(())
});
}
}
fn message_for_offset(&self, offset: usize, cx: &App) -> Option<Message> {
self.messages_for_offsets([offset], cx).pop()
}
pub fn messages_for_offsets(
&self,
offsets: impl IntoIterator<Item = usize>,
cx: &App,
) -> Vec<Message> {
let mut result = Vec::new();
let mut messages = self.messages(cx).peekable();
let mut offsets = offsets.into_iter().peekable();
let mut current_message = messages.next();
while let Some(offset) = offsets.next() {
// Locate the message that contains the offset.
while current_message.as_ref().map_or(false, |message| {
!message.offset_range.contains(&offset) && messages.peek().is_some()
}) {
current_message = messages.next();
}
let Some(message) = current_message.as_ref() else {
break;
};
// Skip offsets that are in the same message.
while offsets.peek().map_or(false, |offset| {
message.offset_range.contains(offset) || messages.peek().is_none()
}) {
offsets.next();
}
result.push(message.clone());
}
result
}
fn messages_from_anchors<'a>(
&'a self,
message_anchors: impl Iterator<Item = &'a MessageAnchor> + 'a,
cx: &'a App,
) -> impl 'a + Iterator<Item = Message> {
let buffer = self.buffer.read(cx);
Self::messages_from_iters(buffer, &self.messages_metadata, message_anchors.enumerate())
}
pub fn messages<'a>(&'a self, cx: &'a App) -> impl 'a + Iterator<Item = Message> {
self.messages_from_anchors(self.message_anchors.iter(), cx)
}
pub fn messages_from_iters<'a>(
buffer: &'a Buffer,
metadata: &'a HashMap<MessageId, MessageMetadata>,
messages: impl Iterator<Item = (usize, &'a MessageAnchor)> + 'a,
) -> impl 'a + Iterator<Item = Message> {
let mut messages = messages.peekable();
iter::from_fn(move || {
if let Some((start_ix, message_anchor)) = messages.next() {
let metadata = metadata.get(&message_anchor.id)?;
let message_start = message_anchor.start.to_offset(buffer);
let mut message_end = None;
let mut end_ix = start_ix;
while let Some((_, next_message)) = messages.peek() {
if next_message.start.is_valid(buffer) {
message_end = Some(next_message.start);
break;
} else {
end_ix += 1;
messages.next();
}
}
let message_end_anchor = message_end.unwrap_or(language::Anchor::MAX);
let message_end = message_end_anchor.to_offset(buffer);
return Some(Message {
index_range: start_ix..end_ix,
offset_range: message_start..message_end,
anchor_range: message_anchor.start..message_end_anchor,
id: message_anchor.id,
role: metadata.role,
status: metadata.status.clone(),
cache: metadata.cache.clone(),
});
}
None
})
}
pub fn save(
&mut self,
debounce: Option<Duration>,
fs: Arc<dyn Fs>,
cx: &mut Context<AssistantContext>,
) {
if self.replica_id() != ReplicaId::default() {
// Prevent saving a remote context for now.
return;
}
self.pending_save = cx.spawn(async move |this, cx| {
if let Some(debounce) = debounce {
cx.background_executor().timer(debounce).await;
}
let (old_path, summary) = this.read_with(cx, |this, _| {
let path = this.path.clone();
let summary = if let Some(summary) = this.summary.content() {
if summary.done {
Some(summary.text.clone())
} else {
None
}
} else {
None
};
(path, summary)
})?;
if let Some(summary) = summary {
let context = this.read_with(cx, |this, cx| this.serialize(cx))?;
let mut discriminant = 1;
let mut new_path;
loop {
new_path = contexts_dir().join(&format!(
"{} - {}.zed.json",
summary.trim(),
discriminant
));
if fs.is_file(&new_path).await {
discriminant += 1;
} else {
break;
}
}
fs.create_dir(contexts_dir().as_ref()).await?;
// rename before write ensures that only one file exists
if let Some(old_path) = old_path.as_ref() {
if new_path.as_path() != old_path.as_ref() {
fs.rename(
&old_path,
&new_path,
RenameOptions {
overwrite: true,
ignore_if_exists: true,
},
)
.await?;
}
}
// update path before write in case it fails
this.update(cx, {
let new_path: Arc<Path> = new_path.clone().into();
move |this, cx| {
this.path = Some(new_path.clone());
cx.emit(ContextEvent::PathChanged { old_path, new_path });
}
})
.ok();
fs.atomic_write(new_path, serde_json::to_string(&context).unwrap())
.await?;
}
Ok(())
});
}
pub fn set_custom_summary(&mut self, custom_summary: String, cx: &mut Context<Self>) {
let timestamp = self.next_timestamp();
let summary = self.summary.content_or_set_empty();
summary.timestamp = timestamp;
summary.done = true;
summary.text = custom_summary;
cx.emit(ContextEvent::SummaryChanged);
}
}
#[derive(Debug, Default)]
pub struct ContextVersion {
context: clock::Global,
buffer: clock::Global,
}
impl ContextVersion {
pub fn from_proto(proto: &proto::ContextVersion) -> Self {
Self {
context: language::proto::deserialize_version(&proto.context_version),
buffer: language::proto::deserialize_version(&proto.buffer_version),
}
}
pub fn to_proto(&self, context_id: ContextId) -> proto::ContextVersion {
proto::ContextVersion {
context_id: context_id.to_proto(),
context_version: language::proto::serialize_version(&self.context),
buffer_version: language::proto::serialize_version(&self.buffer),
}
}
}
#[derive(Debug, Clone)]
pub struct ParsedSlashCommand {
pub name: String,
pub arguments: SmallVec<[String; 3]>,
pub status: PendingSlashCommandStatus,
pub source_range: Range<language::Anchor>,
}
#[derive(Debug)]
pub struct InvokedSlashCommand {
pub name: SharedString,
pub range: Range<language::Anchor>,
pub run_commands_in_ranges: Vec<Range<language::Anchor>>,
pub status: InvokedSlashCommandStatus,
pub transaction: Option<language::TransactionId>,
timestamp: clock::Lamport,
}
#[derive(Debug)]
pub enum InvokedSlashCommandStatus {
Running(Task<()>),
Error(SharedString),
Finished,
}
#[derive(Debug, Clone)]
pub enum PendingSlashCommandStatus {
Idle,
Running { _task: Shared<Task<()>> },
Error(String),
}
#[derive(Debug, Clone)]
pub struct PendingToolUse {
pub id: LanguageModelToolUseId,
pub name: String,
pub input: serde_json::Value,
pub status: PendingToolUseStatus,
pub source_range: Range<language::Anchor>,
}
#[derive(Debug, Clone)]
pub enum PendingToolUseStatus {
Idle,
Running { _task: Shared<Task<()>> },
Error(String),
}
impl PendingToolUseStatus {
pub fn is_idle(&self) -> bool {
matches!(self, PendingToolUseStatus::Idle)
}
}
#[derive(Serialize, Deserialize)]
pub struct SavedMessage {
pub id: MessageId,
pub start: usize,
pub metadata: MessageMetadata,
}
#[derive(Serialize, Deserialize)]
pub struct SavedContext {
pub id: Option<ContextId>,
pub zed: String,
pub version: String,
pub text: String,
pub messages: Vec<SavedMessage>,
pub summary: String,
pub slash_command_output_sections:
Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
#[serde(default)]
pub thought_process_output_sections: Vec<ThoughtProcessOutputSection<usize>>,
}
impl SavedContext {
pub const VERSION: &'static str = "0.4.0";
pub fn from_json(json: &str) -> Result<Self> {
let saved_context_json = serde_json::from_str::<serde_json::Value>(json)?;
match saved_context_json
.get("version")
.context("version not found")?
{
serde_json::Value::String(version) => match version.as_str() {
SavedContext::VERSION => {
Ok(serde_json::from_value::<SavedContext>(saved_context_json)?)
}
SavedContextV0_3_0::VERSION => {
let saved_context =
serde_json::from_value::<SavedContextV0_3_0>(saved_context_json)?;
Ok(saved_context.upgrade())
}
SavedContextV0_2_0::VERSION => {
let saved_context =
serde_json::from_value::<SavedContextV0_2_0>(saved_context_json)?;
Ok(saved_context.upgrade())
}
SavedContextV0_1_0::VERSION => {
let saved_context =
serde_json::from_value::<SavedContextV0_1_0>(saved_context_json)?;
Ok(saved_context.upgrade())
}
_ => anyhow::bail!("unrecognized saved context version: {version:?}"),
},
_ => anyhow::bail!("version not found on saved context"),
}
}
fn into_ops(
self,
buffer: &Entity<Buffer>,
cx: &mut Context<AssistantContext>,
) -> Vec<ContextOperation> {
let mut operations = Vec::new();
let mut version = clock::Global::new();
let mut next_timestamp = clock::Lamport::new(ReplicaId::default());
let mut first_message_metadata = None;
for message in self.messages {
if message.id == MessageId(clock::Lamport::default()) {
first_message_metadata = Some(message.metadata);
} else {
operations.push(ContextOperation::InsertMessage {
anchor: MessageAnchor {
id: message.id,
start: buffer.read(cx).anchor_before(message.start),
},
metadata: MessageMetadata {
role: message.metadata.role,
status: message.metadata.status,
timestamp: message.metadata.timestamp,
cache: None,
},
version: version.clone(),
});
version.observe(message.id.0);
next_timestamp.observe(message.id.0);
}
}
if let Some(metadata) = first_message_metadata {
let timestamp = next_timestamp.tick();
operations.push(ContextOperation::UpdateMessage {
message_id: MessageId(clock::Lamport::default()),
metadata: MessageMetadata {
role: metadata.role,
status: metadata.status,
timestamp,
cache: None,
},
version: version.clone(),
});
version.observe(timestamp);
}
let buffer = buffer.read(cx);
for section in self.slash_command_output_sections {
let timestamp = next_timestamp.tick();
operations.push(ContextOperation::SlashCommandOutputSectionAdded {
timestamp,
section: SlashCommandOutputSection {
range: buffer.anchor_after(section.range.start)
..buffer.anchor_before(section.range.end),
icon: section.icon,
label: section.label,
metadata: section.metadata,
},
version: version.clone(),
});
version.observe(timestamp);
}
for section in self.thought_process_output_sections {
let timestamp = next_timestamp.tick();
operations.push(ContextOperation::ThoughtProcessOutputSectionAdded {
timestamp,
section: ThoughtProcessOutputSection {
range: buffer.anchor_after(section.range.start)
..buffer.anchor_before(section.range.end),
},
version: version.clone(),
});
version.observe(timestamp);
}
let timestamp = next_timestamp.tick();
operations.push(ContextOperation::UpdateSummary {
summary: ContextSummaryContent {
text: self.summary,
done: true,
timestamp,
},
version: version.clone(),
});
version.observe(timestamp);
operations
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
struct SavedMessageIdPreV0_4_0(usize);
#[derive(Serialize, Deserialize)]
struct SavedMessagePreV0_4_0 {
id: SavedMessageIdPreV0_4_0,
start: usize,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
struct SavedMessageMetadataPreV0_4_0 {
role: Role,
status: MessageStatus,
}
#[derive(Serialize, Deserialize)]
struct SavedContextV0_3_0 {
id: Option<ContextId>,
zed: String,
version: String,
text: String,
messages: Vec<SavedMessagePreV0_4_0>,
message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
summary: String,
slash_command_output_sections: Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
}
impl SavedContextV0_3_0 {
const VERSION: &'static str = "0.3.0";
fn upgrade(self) -> SavedContext {
SavedContext {
id: self.id,
zed: self.zed,
version: SavedContext::VERSION.into(),
text: self.text,
messages: self
.messages
.into_iter()
.filter_map(|message| {
let metadata = self.message_metadata.get(&message.id)?;
let timestamp = clock::Lamport {
replica_id: ReplicaId::default(),
value: message.id.0 as u32,
};
Some(SavedMessage {
id: MessageId(timestamp),
start: message.start,
metadata: MessageMetadata {
role: metadata.role,
status: metadata.status.clone(),
timestamp,
cache: None,
},
})
})
.collect(),
summary: self.summary,
slash_command_output_sections: self.slash_command_output_sections,
thought_process_output_sections: Vec::new(),
}
}
}
#[derive(Serialize, Deserialize)]
struct SavedContextV0_2_0 {
id: Option<ContextId>,
zed: String,
version: String,
text: String,
messages: Vec<SavedMessagePreV0_4_0>,
message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
summary: String,
}
impl SavedContextV0_2_0 {
const VERSION: &'static str = "0.2.0";
fn upgrade(self) -> SavedContext {
SavedContextV0_3_0 {
id: self.id,
zed: self.zed,
version: SavedContextV0_3_0::VERSION.to_string(),
text: self.text,
messages: self.messages,
message_metadata: self.message_metadata,
summary: self.summary,
slash_command_output_sections: Vec::new(),
}
.upgrade()
}
}
#[derive(Serialize, Deserialize)]
struct SavedContextV0_1_0 {
id: Option<ContextId>,
zed: String,
version: String,
text: String,
messages: Vec<SavedMessagePreV0_4_0>,
message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
summary: String,
api_url: Option<String>,
model: OpenAiModel,
}
impl SavedContextV0_1_0 {
const VERSION: &'static str = "0.1.0";
fn upgrade(self) -> SavedContext {
SavedContextV0_2_0 {
id: self.id,
zed: self.zed,
version: SavedContextV0_2_0::VERSION.to_string(),
text: self.text,
messages: self.messages,
message_metadata: self.message_metadata,
summary: self.summary,
}
.upgrade()
}
}
#[derive(Debug, Clone)]
pub struct SavedContextMetadata {
pub title: SharedString,
pub path: Arc<Path>,
pub mtime: chrono::DateTime<chrono::Local>,
}
Reference in a new issue View git blame Copy permalink