Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Make thread context wait on detailed summary + remove "Summarizing context..." (#29564)

Browse source 
This moves summarization task management out of `context_store`. The
code there was draining a Vec of tasks to block on, but this is no
longer a good fit for message_editor's context loading. It needs to be
able to repeatedly await on the thread summarization tasks involved in
the context.

Discussed with Danilo, and he thinks it'd be good to remove the current
"Summarizing context" anyway since it causes layout shift. If message
send is blocked on summarizing, the pulsing context pill is sufficient
for now. This UI change made this overall change more straightforward.

Release Notes:

- N/A
This commit is contained in:
Michael Sloan 2025-04-28 15:21:20 -06:00 • committed by GitHub
parent 99df1190a9
commit bbc66748dd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 160 additions and 197 deletions
Download patch file Download diff file Expand all files Collapse all files

135
crates/agent/src/thread.rs
View file

@ -14,7 +14,8 @@ use futures::future::Shared;
use futures::{FutureExt, StreamExt as _};
use git::repository::DiffType;
use gpui::{
AnyWindowHandle, App, AppContext, Context, Entity, EventEmitter, SharedString, Task, WeakEntity,
AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task,
WeakEntity,
};
use language_model::{
ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
@ -24,6 +25,7 @@ use language_model::{
ModelRequestLimitReachedError, PaymentRequiredError, RequestUsage, Role, StopReason,
TokenUsage,
};
use postage::stream::Stream as _;
use project::Project;
use project::git_store::{GitStore, GitStoreCheckpoint, RepositoryState};
use prompt_store::{ModelContext, PromptBuilder};
@ -36,6 +38,7 @@ use util::{ResultExt as _, TryFutureExt as _, post_inc};
use uuid::Uuid;
use zed_llm_client::CompletionMode;
use crate::ThreadStore;
use crate::context::{AgentContext, ContextLoadResult, LoadedContext};
use crate::thread_store::{
SerializedMessage, SerializedMessageSegment, SerializedThread, SerializedToolResult,
@ -243,6 +246,16 @@ pub enum DetailedSummaryState {
},
}
impl DetailedSummaryState {
fn text(&self) -> Option<SharedString> {
if let Self::Generated { text, .. } = self {
Some(text.clone())
} else {
None
}
}
}
#[derive(Default)]
pub struct TotalTokenUsage {
pub total: usize,
@ -290,7 +303,9 @@ pub struct Thread {
updated_at: DateTime<Utc>,
summary: Option<SharedString>,
pending_summary: Task<Option<()>>,
detailed_summary_state: DetailedSummaryState,
detailed_summary_task: Task<Option<()>>,
detailed_summary_tx: postage::watch::Sender<DetailedSummaryState>,
detailed_summary_rx: postage::watch::Receiver<DetailedSummaryState>,
completion_mode: Option<CompletionMode>,
messages: Vec<Message>,
next_message_id: MessageId,
@ -335,12 +350,15 @@ impl Thread {
system_prompt: SharedProjectContext,
cx: &mut Context<Self>,
) -> Self {
let (detailed_summary_tx, detailed_summary_rx) = postage::watch::channel();
Self {
id: ThreadId::new(),
updated_at: Utc::now(),
summary: None,
pending_summary: Task::ready(None),
detailed_summary_state: DetailedSummaryState::NotGenerated,
detailed_summary_task: Task::ready(None),
detailed_summary_tx,
detailed_summary_rx,
completion_mode: None,
messages: Vec::new(),
next_message_id: MessageId(0),
@ -390,13 +408,17 @@ impl Thread {
.unwrap_or(0),
);
let tool_use = ToolUseState::from_serialized_messages(tools.clone(), &serialized.messages);
let (detailed_summary_tx, detailed_summary_rx) =
postage::watch::channel_with(serialized.detailed_summary_state);
Self {
id,
updated_at: serialized.updated_at,
summary: Some(serialized.summary),
pending_summary: Task::ready(None),
detailed_summary_state: serialized.detailed_summary_state,
detailed_summary_task: Task::ready(None),
detailed_summary_tx,
detailed_summary_rx,
completion_mode: None,
messages: serialized
.messages
@ -509,19 +531,6 @@ impl Thread {
}
}
pub fn latest_detailed_summary_or_text(&self) -> SharedString {
self.latest_detailed_summary()
.unwrap_or_else(|| self.text().into())
}
fn latest_detailed_summary(&self) -> Option<SharedString> {
if let DetailedSummaryState::Generated { text, .. } = &self.detailed_summary_state {
Some(text.clone())
} else {
None
}
}
pub fn completion_mode(&self) -> Option<CompletionMode> {
self.completion_mode
}
@ -941,7 +950,7 @@ impl Thread {
initial_project_snapshot,
cumulative_token_usage: this.cumulative_token_usage,
request_token_usage: this.request_token_usage.clone(),
detailed_summary_state: this.detailed_summary_state.clone(),
detailed_summary_state: this.detailed_summary_rx.borrow().clone(),
exceeded_window_error: this.exceeded_window_error.clone(),
})
})
@ -1540,25 +1549,34 @@ impl Thread {
});
}
pub fn generate_detailed_summary(&mut self, cx: &mut Context<Self>) -> Option<Task<()>> {
let last_message_id = self.messages.last().map(|message| message.id)?;
pub fn start_generating_detailed_summary_if_needed(
&mut self,
thread_store: WeakEntity<ThreadStore>,
cx: &mut Context<Self>,
) {
let Some(last_message_id) = self.messages.last().map(|message| message.id) else {
return;
};
match &self.detailed_summary_state {
match &*self.detailed_summary_rx.borrow() {
DetailedSummaryState::Generating { message_id, .. }
| DetailedSummaryState::Generated { message_id, .. }
if *message_id == last_message_id =>
{
// Already up-to-date
return None;
return;
}
_ => {}
}
let ConfiguredModel { model, provider } =
LanguageModelRegistry::read_global(cx).thread_summary_model()?;
let Some(ConfiguredModel { model, provider }) =
LanguageModelRegistry::read_global(cx).thread_summary_model()
else {
return;
};
if !provider.is_authenticated(cx) {
return None;
return;
}
let added_user_message = "Generate a detailed summary of this conversation. Include:\n\
@ -1570,16 +1588,24 @@ impl Thread {
let request = self.to_summarize_request(added_user_message.into());
let task = cx.spawn(async move |thread, cx| {
*self.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generating {
message_id: last_message_id,
};
// Replace the detailed summarization task if there is one, cancelling it. It would probably
// be better to allow the old task to complete, but this would require logic for choosing
// which result to prefer (the old task could complete after the new one, resulting in a
// stale summary).
self.detailed_summary_task = cx.spawn(async move |thread, cx| {
let stream = model.stream_completion_text(request, &cx);
let Some(mut messages) = stream.await.log_err() else {
thread
.update(cx, |this, _cx| {
this.detailed_summary_state = DetailedSummaryState::NotGenerated;
.update(cx, |thread, _cx| {
*thread.detailed_summary_tx.borrow_mut() =
DetailedSummaryState::NotGenerated;
})
.log_err();
return;
.ok()?;
return None;
};
let mut new_detailed_summary = String::new();
@ -1591,25 +1617,56 @@ impl Thread {
}
thread
.update(cx, |this, _cx| {
this.detailed_summary_state = DetailedSummaryState::Generated {
.update(cx, |thread, _cx| {
*thread.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generated {
text: new_detailed_summary.into(),
message_id: last_message_id,
};
})
.log_err();
.ok()?;
// Save thread so its summary can be reused later
if let Some(thread) = thread.upgrade() {
if let Ok(Ok(save_task)) = cx.update(|cx| {
thread_store
.update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx))
}) {
save_task.await.log_err();
}
}
Some(())
});
}
self.detailed_summary_state = DetailedSummaryState::Generating {
message_id: last_message_id,
};
pub async fn wait_for_detailed_summary_or_text(
this: &Entity<Self>,
cx: &mut AsyncApp,
) -> Option<SharedString> {
let mut detailed_summary_rx = this
.read_with(cx, |this, _cx| this.detailed_summary_rx.clone())
.ok()?;
loop {
match detailed_summary_rx.recv().await? {
DetailedSummaryState::Generating { .. } => {}
DetailedSummaryState::NotGenerated => {
return this.read_with(cx, |this, _cx| this.text().into()).ok();
}
DetailedSummaryState::Generated { text, .. } => return Some(text),
}
}
}
Some(task)
pub fn latest_detailed_summary_or_text(&self) -> SharedString {
self.detailed_summary_rx
.borrow()
.text()
.unwrap_or_else(|| self.text().into())
}
pub fn is_generating_detailed_summary(&self) -> bool {
matches!(
self.detailed_summary_state,
&*self.detailed_summary_rx.borrow(),
DetailedSummaryState::Generating { .. }
)
}