From e9b18cd2d54153101dda72aa59f8b4dea1f52372 Mon Sep 17 00:00:00 2001 From: Prohect Date: Wed, 26 Aug 2026 14:33:17 +0000 Subject: [PATCH] agent: Coalesce thread saves into a single worker (#63053) # Objective - Reduce excessive RAM usage during long agent sessions. Previously every `Thread::notify` triggered a full serialization of the entire conversation (including any embedded base64 image) to JSON, so streaming bursts queued many large, redundant saves and ballooned memory into the multi-GB range. - Related to #57126. - Related to #62133. ## Solution - Replace the per-notification save task with a single coalescing save worker per session in `crates/agent/src/agent.rs`. - `save_thread` now writes the latest snapshot into an in-place buffer and wakes the worker via a `watch` channel, instead of spawning a fresh full-conversation save on every notification. - The worker drains the buffer in a loop, saving only the latest snapshot per burst (latest-wins coalescing), and preserves save-on-close by draining the final snapshot before exiting. ## Testing - `cargo check -p agent` - `cargo test -p agent test_close_session_saves_thread` - `cargo test -p agent test_save_load_thread` - `cargo test -p agent test_threads_flushed_to_database_on_app_quit` ## Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content adheres to Zed's UI standards (no UI change) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable --- Release Notes: - Improved memory usage during long agent sessions. --------- Co-authored-by: Kirill Bulatov --- crates/agent/src/agent.rs | 199 +++++++++++++++++++++++++++++++++++--- crates/agent/src/db.rs | 24 ++++- 2 files changed, 205 insertions(+), 18 deletions(-) diff --git a/crates/agent/src/agent.rs b/crates/agent/src/agent.rs index 1b46bdfb25..6c8537c294 100644 --- a/crates/agent/src/agent.rs +++ b/crates/agent/src/agent.rs @@ -16,6 +16,7 @@ use context_server::ContextServerId; pub use db::*; use itertools::Itertools; pub use native_agent_server::NativeAgentServer; +use parking_lot::Mutex; pub use pattern_extraction::*; pub use sandboxing::{ ThreadSandbox, sandbox_worktree_writable_paths, settings_sandbox_policy, @@ -214,6 +215,14 @@ struct ProjectState { _subscriptions: Vec, } +/// A thread snapshot captured for persistence. Building it clones the thread's +/// messages (cheap `Arc` clones) so the background save can serialize without +/// racing the foreground mutation of the live thread. +struct PendingThreadSave { + folder_paths: PathList, + db_thread: Task, +} + /// Holds both the internal Thread and the AcpThread for a session struct Session { /// The internal thread that processes messages @@ -221,7 +230,11 @@ struct Session { /// The ACP thread that handles protocol communication acp_thread: Entity, project_id: EntityId, - pending_save: Task>, + /// 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>>, + save_wake: watch::Sender<()>, + save_worker: Task>, _subscriptions: Vec, ref_count: usize, } @@ -837,14 +850,36 @@ impl NativeAgent { }), ]; + let (save_wake, save_wake_rx) = watch::channel(()); + let pending_save: Arc>> = Arc::new(Mutex::new(None)); + let database_future = ThreadsDatabase::connect(cx); + let thread_store = self.thread_store.clone(); + let save_worker = cx.spawn({ + let pending_save = pending_save.clone(); + let session_id = session_id.clone(); + async move |_this, cx| { + Self::run_save_worker( + session_id, + save_wake_rx, + pending_save, + database_future, + thread_store, + cx, + ) + .await + } + }); + self.sessions.insert( session_id, Session { thread: thread_handle, acp_thread: acp_thread.clone(), project_id, + pending_save, + save_wake, + save_worker, _subscriptions: subscriptions, - pending_save: Task::ready(Ok(())), ref_count, }, ); @@ -1745,7 +1780,7 @@ impl NativeAgent { self.publish_skill_index(cx); } - session.pending_save + session.save_worker } fn save_thread(&mut self, thread: Entity, cx: &mut Context) { @@ -1757,23 +1792,49 @@ impl NativeAgent { return; }; - let database_future = ThreadsDatabase::connect(cx); - let thread_store = self.thread_store.clone(); let Some(session) = self.sessions.get_mut(&id) else { return; }; - session.pending_save = cx.spawn(async move |_, cx| { - let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else { - return Ok(()); - }; - let db_thread = db_thread.await; - database - .save_thread(id, db_thread, folder_paths) - .await - .log_err(); - thread_store.update(cx, |store, cx| store.reload(cx)); - Ok(()) + *session.pending_save.lock() = Some(PendingThreadSave { + folder_paths, + db_thread, }); + session.save_wake.send(()).log_err(); + } + + async fn run_save_worker( + id: acp::SessionId, + mut wake: watch::Receiver<()>, + pending_save: Arc>>, + database_future: Shared, Arc>>>, + thread_store: Entity, + cx: &mut AsyncApp, + ) -> Result<()> { + loop { + let closed = wake.changed().await.is_err(); + let payload = pending_save.lock().take(); + if let Some(PendingThreadSave { + folder_paths, + db_thread, + }) = payload + && let Some(database) = database_future + .clone() + .await + .map_err(|err| anyhow!(err)) + .log_err() + { + let db_thread = db_thread.await; + database + .save_thread(id.clone(), db_thread, folder_paths) + .await + .log_err(); + thread_store.update(cx, |store, cx| store.reload(cx)); + } + if closed { + break; + } + } + Ok(()) } /// Builds everything needed to persist a session's thread content, @@ -6600,6 +6661,97 @@ mod internal_tests { }); } + #[gpui::test] + async fn test_save_burst_preserves_in_flight_write(cx: &mut TestAppContext) { + init_test(cx); + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + "/", + json!({ + "a": { + "file.txt": "hello" + } + }), + ) + .await; + let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await; + let thread_store = cx.new(|cx| ThreadStore::new(cx)); + let agent = cx + .update(|cx| NativeAgent::new(thread_store.clone(), Templates::new(), fs.clone(), cx)); + let connection = Rc::new(NativeAgentConnection(agent.clone())); + + let acp_thread = cx + .update(|cx| { + connection + .clone() + .new_session(project.clone(), PathList::new(&[Path::new("")]), cx) + }) + .await + .unwrap(); + let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone()); + let thread = agent.read_with(cx, |agent, _| { + agent.sessions.get(&session_id).unwrap().thread.clone() + }); + + let model = Arc::new(FakeLanguageModel::default()); + thread.update(cx, |thread, cx| { + thread.set_model(model.clone(), cx); + }); + + 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(); + + model.send_last_completion_stream_text_chunk("world"); + model.end_last_completion_stream(); + send.await.unwrap(); + cx.run_until_parked(); + + let database = cx.update(|cx| ThreadsDatabase::connect(cx)).await.unwrap(); + + let [first_draft, second_draft, third_draft] = ["draft one", "draft two", "draft three"] + .map(|text| vec![acp::ContentBlock::Text(acp::TextContent::new(text))]); + + let (first_gate_tx, first_gate_rx) = oneshot::channel(); + database.set_write_gate(first_gate_rx); + set_draft_and_save(&agent, &acp_thread, &thread, first_draft.clone(), cx); + cx.run_until_parked(); + + set_draft_and_save(&agent, &acp_thread, &thread, second_draft, cx); + set_draft_and_save(&agent, &acp_thread, &thread, third_draft.clone(), cx); + cx.run_until_parked(); + + let (second_gate_tx, second_gate_rx) = oneshot::channel(); + database.set_write_gate(second_gate_rx); + first_gate_tx.send(()).ok(); + cx.run_until_parked(); + + let db_thread = database + .load_thread(session_id.clone()) + .await + .unwrap() + .unwrap(); + assert_eq!( + db_thread.draft_prompt, + Some(first_draft), + "save requests arriving during an in-flight write must not cancel it" + ); + + second_gate_tx.send(()).ok(); + cx.run_until_parked(); + + let db_thread = database + .load_thread(session_id.clone()) + .await + .unwrap() + .unwrap(); + assert_eq!( + db_thread.draft_prompt, + Some(third_draft), + "the save worker must persist the latest snapshot of a burst" + ); + } + #[gpui::test] async fn test_thread_summary_releases_loaded_session(cx: &mut TestAppContext) { init_test(cx); @@ -6958,6 +7110,21 @@ mod internal_tests { // silently mangling unrelated user text. assert_eq!(strip_slash_command_prefix("hello world"), "hello world",); } + + fn set_draft_and_save( + agent: &Entity, + acp_thread: &Entity, + thread: &Entity, + draft: Vec, + cx: &mut TestAppContext, + ) { + acp_thread.update(cx, |acp_thread, cx| { + acp_thread.set_draft_prompt(Some(draft), cx); + }); + agent.update(cx, |agent, cx| { + agent.save_thread(thread.clone(), cx); + }); + } } fn mcp_message_content_to_acp_content_block( diff --git a/crates/agent/src/db.rs b/crates/agent/src/db.rs index 0dd91bfdea..e7e9bcdc47 100644 --- a/crates/agent/src/db.rs +++ b/crates/agent/src/db.rs @@ -392,6 +392,12 @@ impl Column for DataType { pub(crate) struct ThreadsDatabase { executor: BackgroundExecutor, connection: Arc>, + /// In production, saves take real time (serialization, zstd, disk I/O) while + /// the user keeps typing, so new save requests routinely arrive mid-write. + /// The test executor completes writes instantly, so tests use this gate to + /// hold a write in flight and interleave more save requests with it. + #[cfg(test)] + write_gate: Mutex>>>, } struct GlobalThreadsDatabase(Shared, Arc>>>); @@ -481,6 +487,8 @@ impl ThreadsDatabase { let db = Self { executor, connection: Arc::new(Mutex::new(connection)), + #[cfg(test)] + write_gate: Mutex::new(None), }; Ok(db) @@ -629,9 +637,21 @@ impl ThreadsDatabase { folder_paths: PathList, ) -> Task> { let connection = self.connection.clone(); + #[cfg(test)] + let write_gate = self.write_gate.lock().clone(); - self.executor - .spawn(async move { Self::save_thread_sync(&connection, id, thread, &folder_paths) }) + self.executor.spawn(async move { + #[cfg(test)] + if let Some(write_gate) = write_gate { + write_gate.await.ok(); + } + Self::save_thread_sync(&connection, id, thread, &folder_paths) + }) + } + + #[cfg(test)] + pub fn set_write_gate(&self, gate: futures::channel::oneshot::Receiver<()>) { + *self.write_gate.lock() = Some(gate.shared()); } fn deserialize_thread(data_type: DataType, data: Vec) -> Result {