mirror of
https://github.com/zed-industries/zed.git
synced 2026-10-20 11:46:49 +00:00
agent: Skip redundant thread saves while a response streams (#65244)
Closes #64823 While the agent replied, every streamed chunk saved the whole thread: serializing it to JSON, compressing it, and rewriting the SQLite row. The streamed text isn't saved until its message completes, so nearly all of these writes were identical. Now the thread is not saved while a message is streaming, unless a message has completed or a setting, the title, or the draft prompt changed since the last save. Measured on a 16.6 MB thread (0.76 MB compressed) while the agent replied to one prompt for ~40s: | | Before | After | |---|---|---| | Full-thread saves during the reply | 836 (~21/s) | 2 | | Zed's memory (Activity Monitor) while streaming | 1.34 GB → ~2.2 GB | stays flat | | CPU while streaming | ~75% of a core | ~35% of a core | Before: Initial: <img width="1424" height="200" alt="CleanShot 2026-10-06 at 8 36 34 PM@2x" src="https://github.com/user-attachments/assets/b3065bca-a881-4ff5-9eda-92ced2659554" /> Streaming: <img width="1430" height="212" alt="CleanShot 2026-10-06 at 8 37 20 PM@2x" src="https://github.com/user-attachments/assets/589842b9-7893-4750-8911-4362e981de7c" /> After: Initial: <img width="1358" height="182" alt="CleanShot 2026-10-06 at 8 41 22 PM@2x" src="https://github.com/user-attachments/assets/ea55e40e-4d8c-4218-bf52-6285d7285963" /> Streaming: <img width="1352" height="194" alt="CleanShot 2026-10-06 at 8 38 20 PM@2x" src="https://github.com/user-attachments/assets/616ccede-6303-4da6-96e2-29803aad28aa" /> Release Notes: - Improved memory and CPU usage while the agent responds in long threads.
This commit is contained in:
Smit Barmase
GitHub
parent
8323e23277
commit
98e168265b
4 changed files with 361 additions and 0 deletions
|
|
@ -3398,6 +3398,8 @@ pub struct AcpThread {
|
|||
had_error: bool,
|
||||
/// The user's unsent prompt text, persisted so it can be restored when reloading the thread.
|
||||
draft_prompt: Option<Vec<acp_v2::ContentBlock>>,
|
||||
/// Lets observers detect draft changes without comparing prompts.
|
||||
draft_prompt_revision: usize,
|
||||
/// The initial scroll position for the thread view, set during session registration.
|
||||
ui_scroll_position: Option<gpui::ListOffset>,
|
||||
/// A cursor over retained source, rather than a second text store, lets the UI
|
||||
|
|
@ -3733,6 +3735,7 @@ impl AcpThread {
|
|||
pending_terminal_exit: HashMap::default(),
|
||||
had_error: false,
|
||||
draft_prompt: None,
|
||||
draft_prompt_revision: 0,
|
||||
ui_scroll_position: None,
|
||||
streaming_text_buffer: None,
|
||||
idle_sleep_prevention: IdleSleepPrevention::Inactive,
|
||||
|
|
@ -3768,6 +3771,10 @@ impl AcpThread {
|
|||
self.draft_prompt.as_deref()
|
||||
}
|
||||
|
||||
pub fn draft_prompt_revision(&self) -> usize {
|
||||
self.draft_prompt_revision
|
||||
}
|
||||
|
||||
pub fn set_draft_prompt(
|
||||
&mut self,
|
||||
prompt: Option<Vec<acp_v2::ContentBlock>>,
|
||||
|
|
@ -3775,6 +3782,7 @@ impl AcpThread {
|
|||
) {
|
||||
cx.emit(AcpThreadEvent::PromptUpdated);
|
||||
self.draft_prompt = prompt;
|
||||
self.draft_prompt_revision += 1;
|
||||
}
|
||||
|
||||
pub fn ui_scroll_position(&self) -> Option<gpui::ListOffset> {
|
||||
|
|
|
|||
|
|
@ -235,6 +235,8 @@ struct Session {
|
|||
/// Latest snapshot to persist. Overwritten in place on every save request;
|
||||
/// the single save worker drains it, coalescing bursts into one write.
|
||||
pending_save: Arc<Mutex<Option<PendingThreadSave>>>,
|
||||
last_streaming_save_key: Option<StreamingSaveKey>,
|
||||
last_draft_prompt_revision: Option<usize>,
|
||||
save_wake: watch::Sender<()>,
|
||||
save_worker: Task<Result<()>>,
|
||||
_subscriptions: Vec<Subscription>,
|
||||
|
|
@ -247,6 +249,24 @@ impl Session {
|
|||
None => self.thread.read(cx).draft_prompt().map(Vec::from),
|
||||
}
|
||||
}
|
||||
|
||||
fn draft_prompt_revision(&self, cx: &App) -> Option<usize> {
|
||||
self.acp_thread
|
||||
.upgrade()
|
||||
.map(|acp_thread| acp_thread.read(cx).draft_prompt_revision())
|
||||
}
|
||||
|
||||
fn can_skip_save_while_streaming(&self, cx: &App) -> bool {
|
||||
// Streaming notifies the thread once per chunk, but `to_db` doesn't
|
||||
// include the streaming message, so those notifies have nothing new to save.
|
||||
let thread = self.thread.read(cx);
|
||||
thread.is_streaming_message()
|
||||
&& self
|
||||
.last_streaming_save_key
|
||||
.as_ref()
|
||||
.is_some_and(|key| *key == thread.streaming_save_key())
|
||||
&& self.last_draft_prompt_revision == self.draft_prompt_revision(cx)
|
||||
}
|
||||
}
|
||||
|
||||
struct PendingSession {
|
||||
|
|
@ -894,6 +914,8 @@ impl NativeAgent {
|
|||
subagents: Vec::new(),
|
||||
project_id,
|
||||
pending_save,
|
||||
last_streaming_save_key: None,
|
||||
last_draft_prompt_revision: None,
|
||||
save_wake,
|
||||
save_worker,
|
||||
_subscriptions: subscriptions,
|
||||
|
|
@ -1804,6 +1826,9 @@ impl NativeAgent {
|
|||
let Some(session) = self.sessions.get(&id) else {
|
||||
return;
|
||||
};
|
||||
if session.can_skip_save_while_streaming(cx) {
|
||||
return;
|
||||
}
|
||||
let draft_prompt = session.draft_prompt(cx);
|
||||
self.enqueue_save(&id, draft_prompt, cx);
|
||||
}
|
||||
|
|
@ -1826,6 +1851,8 @@ impl NativeAgent {
|
|||
let Some(session) = self.sessions.get_mut(&id) else {
|
||||
return;
|
||||
};
|
||||
session.last_streaming_save_key = Some(session.thread.read(cx).streaming_save_key());
|
||||
session.last_draft_prompt_revision = session.draft_prompt_revision(cx);
|
||||
*session.pending_save.lock() = Some(PendingThreadSave {
|
||||
folder_paths,
|
||||
db_thread,
|
||||
|
|
@ -7669,6 +7696,279 @@ mod internal_tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_streaming_chunks_do_not_save_thread(cx: &mut TestAppContext) {
|
||||
let (turn, send) = StreamingTurn::start(cx).await;
|
||||
let saves_before_stream = turn.database.save_count();
|
||||
|
||||
for chunk in ["one ", "two ", "three"] {
|
||||
turn.fake.send_last_text(&turn.model, chunk);
|
||||
cx.run_until_parked();
|
||||
}
|
||||
assert_eq!(
|
||||
turn.database.save_count(),
|
||||
saves_before_stream,
|
||||
"streamed chunks must not save the thread"
|
||||
);
|
||||
|
||||
turn.fake.end_last(&turn.model);
|
||||
send.await.unwrap();
|
||||
cx.run_until_parked();
|
||||
assert!(turn.database.save_count() > saves_before_stream);
|
||||
assert_eq!(
|
||||
turn.saved_markdown().await,
|
||||
"## User\n\nhello\n\n## Assistant\n\none two three\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_completed_message_saved_while_next_message_streams(cx: &mut TestAppContext) {
|
||||
let (turn, send) = StreamingTurn::start(cx).await;
|
||||
|
||||
turn.fake.send_last_event(
|
||||
&turn.model,
|
||||
LanguageModelCompletionEvent::StartMessage {
|
||||
message_id: "first".into(),
|
||||
},
|
||||
);
|
||||
turn.fake.send_last_text(&turn.model, "first message");
|
||||
turn.fake.send_last_event(
|
||||
&turn.model,
|
||||
LanguageModelCompletionEvent::StartMessage {
|
||||
message_id: "second".into(),
|
||||
},
|
||||
);
|
||||
turn.fake.send_last_text(&turn.model, "second mess");
|
||||
cx.run_until_parked();
|
||||
assert_eq!(
|
||||
turn.saved_markdown().await,
|
||||
"## User\n\nhello\n\n## Assistant\n\nfirst message\n",
|
||||
"a completed message must be saved even while the next one streams"
|
||||
);
|
||||
|
||||
turn.fake.send_last_text(&turn.model, "age");
|
||||
turn.fake.end_last(&turn.model);
|
||||
send.await.unwrap();
|
||||
cx.run_until_parked();
|
||||
assert_eq!(
|
||||
turn.saved_markdown().await,
|
||||
"## User\n\nhello\n\n## Assistant\n\nfirst message\n\n## Assistant\n\nsecond message\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_cancelling_stream_saves_partial_message(cx: &mut TestAppContext) {
|
||||
let (turn, send) = StreamingTurn::start(cx).await;
|
||||
|
||||
turn.fake.send_last_text(&turn.model, "partial");
|
||||
cx.run_until_parked();
|
||||
turn.acp_thread
|
||||
.update(cx, |thread, cx| thread.cancel(cx))
|
||||
.await;
|
||||
send.await.unwrap();
|
||||
cx.run_until_parked();
|
||||
assert_eq!(
|
||||
turn.saved_markdown().await,
|
||||
"## User\n\nhello\n\n## Assistant\n\npartial\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_settings_changes_save_while_message_streams(cx: &mut TestAppContext) {
|
||||
let (turn, send) = StreamingTurn::start(cx).await;
|
||||
turn.fake.send_last_text(&turn.model, "one ");
|
||||
cx.run_until_parked();
|
||||
|
||||
turn.thread.update(cx, |thread, cx| {
|
||||
thread.set_thinking_effort(Some("high".into()), cx);
|
||||
});
|
||||
cx.run_until_parked();
|
||||
assert_eq!(
|
||||
turn.saved_thread().await.thinking_effort.as_deref(),
|
||||
Some("high")
|
||||
);
|
||||
|
||||
turn.thread
|
||||
.update(cx, |thread, cx| thread.set_title("Renamed".into(), cx));
|
||||
cx.run_until_parked();
|
||||
assert_eq!(turn.saved_thread().await.title, "Renamed");
|
||||
|
||||
let draft = vec![acp_v2::ContentBlock::Text(acp_v2::TextContent::new(
|
||||
"next prompt",
|
||||
))];
|
||||
turn.acp_thread.update(cx, |acp_thread, cx| {
|
||||
acp_thread.set_draft_prompt(Some(draft.clone()), cx);
|
||||
});
|
||||
// The thread view notifies the thread after a draft edit to request a save.
|
||||
turn.thread.update(cx, |_, cx| cx.notify());
|
||||
cx.run_until_parked();
|
||||
assert_eq!(turn.saved_thread().await.draft_prompt, Some(draft));
|
||||
|
||||
let saves_after_changes = turn.database.save_count();
|
||||
turn.fake.send_last_text(&turn.model, "two");
|
||||
cx.run_until_parked();
|
||||
assert_eq!(
|
||||
turn.database.save_count(),
|
||||
saves_after_changes,
|
||||
"streamed chunks after a saved change must not save the thread"
|
||||
);
|
||||
|
||||
turn.fake.end_last(&turn.model);
|
||||
send.await.unwrap();
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_token_usage_does_not_save_while_message_streams(cx: &mut TestAppContext) {
|
||||
let (turn, send) = StreamingTurn::start(cx).await;
|
||||
turn.fake.send_last_text(&turn.model, "one ");
|
||||
cx.run_until_parked();
|
||||
let saves_before_usage = turn.database.save_count();
|
||||
|
||||
// Some providers, like Google, report usage with every streamed chunk.
|
||||
for output_tokens in [10, 20, 30] {
|
||||
turn.fake.send_last_event(
|
||||
&turn.model,
|
||||
LanguageModelCompletionEvent::UsageUpdate(language_model::TokenUsage {
|
||||
input_tokens: 100,
|
||||
output_tokens,
|
||||
..Default::default()
|
||||
}),
|
||||
);
|
||||
cx.run_until_parked();
|
||||
}
|
||||
assert_eq!(turn.database.save_count(), saves_before_usage);
|
||||
|
||||
turn.fake.end_last(&turn.model);
|
||||
send.await.unwrap();
|
||||
cx.run_until_parked();
|
||||
let expected_usage = language_model::TokenUsage {
|
||||
input_tokens: 100,
|
||||
output_tokens: 30,
|
||||
..Default::default()
|
||||
};
|
||||
let saved_thread = turn.saved_thread().await;
|
||||
assert_eq!(saved_thread.cumulative_token_usage, expected_usage);
|
||||
assert_eq!(
|
||||
saved_thread
|
||||
.request_token_usage
|
||||
.into_values()
|
||||
.collect::<Vec<_>>(),
|
||||
vec![expected_usage]
|
||||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_token_usage_saved_when_message_ends_without_content(cx: &mut TestAppContext) {
|
||||
let (turn, send) = StreamingTurn::start(cx).await;
|
||||
// A titled thread doesn't generate a title when the turn ends, which
|
||||
// would otherwise save the thread.
|
||||
turn.thread
|
||||
.update(cx, |thread, cx| thread.set_title("Title".into(), cx));
|
||||
turn.send_start_message_with_usage(100);
|
||||
cx.run_until_parked();
|
||||
|
||||
turn.fake.end_last(&turn.model);
|
||||
send.await.unwrap();
|
||||
cx.run_until_parked();
|
||||
assert_eq!(
|
||||
turn.saved_thread()
|
||||
.await
|
||||
.cumulative_token_usage
|
||||
.input_tokens,
|
||||
100
|
||||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_token_usage_saved_when_cancelled_before_content(cx: &mut TestAppContext) {
|
||||
let (turn, send) = StreamingTurn::start(cx).await;
|
||||
// A titled thread doesn't generate a title when the turn ends, which
|
||||
// would otherwise save the thread.
|
||||
turn.thread
|
||||
.update(cx, |thread, cx| thread.set_title("Title".into(), cx));
|
||||
turn.send_start_message_with_usage(100);
|
||||
cx.run_until_parked();
|
||||
|
||||
turn.acp_thread
|
||||
.update(cx, |thread, cx| thread.cancel(cx))
|
||||
.await;
|
||||
send.await.unwrap();
|
||||
cx.run_until_parked();
|
||||
assert_eq!(
|
||||
turn.saved_thread()
|
||||
.await
|
||||
.cumulative_token_usage
|
||||
.input_tokens,
|
||||
100
|
||||
);
|
||||
}
|
||||
|
||||
struct StreamingTurn {
|
||||
fake: Arc<FakeLanguageModelProvider>,
|
||||
model: LanguageModel,
|
||||
acp_thread: Entity<AcpThread>,
|
||||
thread: Entity<Thread>,
|
||||
session_id: acp_v1::SessionId,
|
||||
database: Arc<ThreadsDatabase>,
|
||||
}
|
||||
|
||||
impl StreamingTurn {
|
||||
async fn start(
|
||||
cx: &mut TestAppContext,
|
||||
) -> (Self, Task<Result<Option<acp_thread::SubmissionResponse>>>) {
|
||||
let fake = init_test(cx);
|
||||
let (_connection, agent, _project, acp_thread) = setup_native_agent_session(cx).await;
|
||||
let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
|
||||
let thread = cx.update(|cx| native_thread_for_session(&agent, &session_id, cx));
|
||||
let model = fake.model("fake");
|
||||
thread.update(cx, |thread, cx| thread.set_model(model.clone(), cx));
|
||||
let database = cx.update(|cx| ThreadsDatabase::connect(cx)).await.unwrap();
|
||||
|
||||
let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx));
|
||||
let send = cx.foreground_executor().spawn(send);
|
||||
cx.run_until_parked();
|
||||
|
||||
let turn = Self {
|
||||
fake,
|
||||
model,
|
||||
acp_thread,
|
||||
thread,
|
||||
session_id,
|
||||
database,
|
||||
};
|
||||
(turn, send)
|
||||
}
|
||||
|
||||
async fn saved_thread(&self) -> DbThread {
|
||||
self.database
|
||||
.load_thread(self.session_id.clone())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn saved_markdown(&self) -> String {
|
||||
crate::thread::messages_to_markdown(&self.saved_thread().await.messages)
|
||||
}
|
||||
|
||||
/// Providers like Anthropic report usage before streaming any content.
|
||||
fn send_start_message_with_usage(&self, input_tokens: u64) {
|
||||
self.fake.send_last_event(
|
||||
&self.model,
|
||||
LanguageModelCompletionEvent::StartMessage {
|
||||
message_id: "message".into(),
|
||||
},
|
||||
);
|
||||
self.fake.send_last_event(
|
||||
&self.model,
|
||||
LanguageModelCompletionEvent::UsageUpdate(language_model::TokenUsage {
|
||||
input_tokens,
|
||||
..Default::default()
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_thread_summary_releases_loaded_session(cx: &mut TestAppContext) {
|
||||
let fake = init_test(cx);
|
||||
|
|
|
|||
|
|
@ -399,6 +399,8 @@ pub(crate) struct ThreadsDatabase {
|
|||
/// hold a write in flight and interleave more save requests with it.
|
||||
#[cfg(test)]
|
||||
write_gate: Mutex<Option<Shared<futures::channel::oneshot::Receiver<()>>>>,
|
||||
#[cfg(test)]
|
||||
save_count: std::sync::atomic::AtomicUsize,
|
||||
}
|
||||
|
||||
struct GlobalThreadsDatabase(Shared<Task<Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>>);
|
||||
|
|
@ -490,6 +492,8 @@ impl ThreadsDatabase {
|
|||
connection: Arc::new(Mutex::new(connection)),
|
||||
#[cfg(test)]
|
||||
write_gate: Mutex::new(None),
|
||||
#[cfg(test)]
|
||||
save_count: Default::default(),
|
||||
};
|
||||
|
||||
Ok(db)
|
||||
|
|
@ -640,6 +644,9 @@ impl ThreadsDatabase {
|
|||
let connection = self.connection.clone();
|
||||
#[cfg(test)]
|
||||
let write_gate = self.write_gate.lock().clone();
|
||||
#[cfg(test)]
|
||||
self.save_count
|
||||
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
self.executor.spawn(async move {
|
||||
#[cfg(test)]
|
||||
|
|
@ -650,6 +657,11 @@ impl ThreadsDatabase {
|
|||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn save_count(&self) -> usize {
|
||||
self.save_count.load(std::sync::atomic::Ordering::SeqCst)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn set_write_gate(&self, gate: futures::channel::oneshot::Receiver<()>) {
|
||||
*self.write_gate.lock() = Some(gate.shared());
|
||||
|
|
|
|||
|
|
@ -1270,6 +1270,23 @@ impl From<&ThreadModel> for Option<DbLanguageModel> {
|
|||
}
|
||||
}
|
||||
|
||||
/// The parts of `to_db` that are worth saving while a message streams. Token
|
||||
/// usage and scroll position are left out because they can change on every
|
||||
/// streamed chunk or scroll, and some fields never change after creation.
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct StreamingSaveKey {
|
||||
message_count: usize,
|
||||
title: Option<SharedString>,
|
||||
summary: Option<SharedString>,
|
||||
model: Option<DbLanguageModel>,
|
||||
profile_id: AgentProfileId,
|
||||
speed: Option<Speed>,
|
||||
thinking_enabled: bool,
|
||||
thinking_effort: Option<String>,
|
||||
sandboxed_terminal_temp_dir: Option<PathBuf>,
|
||||
sandbox_grants: crate::db::DbSandboxGrants,
|
||||
}
|
||||
|
||||
pub struct Thread {
|
||||
id: acp::SessionId,
|
||||
prompt_id: PromptId,
|
||||
|
|
@ -1925,6 +1942,8 @@ impl Thread {
|
|||
crate::sandboxing::sandbox_worktree_writable_paths(self.project.read(cx), cx)
|
||||
}
|
||||
|
||||
/// A field added here must also go in `StreamingSaveKey`, unless saving it
|
||||
/// can wait until the response finishes streaming.
|
||||
pub fn to_db(&self, cx: &App) -> Task<DbThread> {
|
||||
let initial_project_snapshot = self.initial_project_snapshot.clone();
|
||||
let mut thread = DbThread {
|
||||
|
|
@ -1959,6 +1978,25 @@ impl Thread {
|
|||
})
|
||||
}
|
||||
|
||||
pub(crate) fn is_streaming_message(&self) -> bool {
|
||||
self.pending_message.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn streaming_save_key(&self) -> StreamingSaveKey {
|
||||
StreamingSaveKey {
|
||||
message_count: self.messages.len(),
|
||||
title: self.title.clone(),
|
||||
summary: self.summary.clone(),
|
||||
model: (&self.model).into(),
|
||||
profile_id: self.profile_id.clone(),
|
||||
speed: self.speed,
|
||||
thinking_enabled: self.thinking_enabled,
|
||||
thinking_effort: self.thinking_effort.clone(),
|
||||
sandboxed_terminal_temp_dir: self.sandboxed_terminal_temp_dir.clone(),
|
||||
sandbox_grants: self.sandbox_grants.borrow().to_db(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a snapshot of the current project state including git information and unsaved buffers.
|
||||
fn project_snapshot(
|
||||
project: Entity<Project>,
|
||||
|
|
@ -4079,6 +4117,9 @@ impl Thread {
|
|||
};
|
||||
|
||||
if message.content.is_empty() {
|
||||
// Saves are skipped while a message streams, so notify to save
|
||||
// anything that changed meanwhile, like token usage.
|
||||
cx.notify();
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue