diff --git a/crates/agent/src/agent.rs b/crates/agent/src/agent.rs index 148e60b2bf8..dff6744cdfc 100644 --- a/crates/agent/src/agent.rs +++ b/crates/agent/src/agent.rs @@ -553,10 +553,15 @@ impl NativeAgent { log::debug!("Creating new NativeAgent"); cx.new(|cx| { - let subscriptions = vec![cx.subscribe( - &LanguageModelRegistry::global(cx), - Self::handle_models_updated_event, - )]; + let subscriptions = vec![ + cx.subscribe( + &LanguageModelRegistry::global(cx), + Self::handle_models_updated_event, + ), + // Flush thread content on quit so an in-flight async save + // can't leave a thread orphaned ("no thread found with ID"). + cx.on_app_quit(Self::flush_threads_on_quit), + ]; if !cx.has_global::() { cx.set_global(SkillIndex::default()); @@ -1765,6 +1770,48 @@ impl NativeAgent { }); } + /// Commits every non-empty thread's content on shutdown so the async + /// `save_thread` losing the race can't leave metadata without content. + fn flush_threads_on_quit( + &mut self, + cx: &mut Context, + ) -> impl Future + use<> { + let database_future = ThreadsDatabase::connect(cx); + + let mut saves = Vec::new(); + for session in self.sessions.values() { + let thread = session.thread.read(cx); + if thread.is_empty() { + continue; + } + let Some(state) = self.projects.get(&session.project_id) else { + continue; + }; + let folder_paths = PathList::new( + &state + .project + .read(cx) + .visible_worktrees(cx) + .map(|worktree| worktree.read(cx).abs_path().to_path_buf()) + .collect::>(), + ); + saves.push((thread.id().clone(), folder_paths, thread.to_db(cx))); + } + + async move { + let Ok(database) = database_future.await else { + return; + }; + for (id, folder_paths, db_thread) in saves { + let db_thread = db_thread.await; + database + .save_thread(id, db_thread, folder_paths) + .await + .log_err(); + } + } + } + fn send_mcp_prompt( &self, message_id: UserMessageId, @@ -3749,6 +3796,58 @@ mod internal_tests { prompt_task.await.unwrap(); } + #[gpui::test] + async fn test_threads_flushed_to_database_on_app_quit(cx: &mut TestAppContext) { + 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)); + + // Give the thread content so it's no longer an empty draft. + cx.update(|cx| { + let path_style = project.read(cx).path_style(cx); + thread.update(cx, |thread, cx| { + thread.push_acp_user_block( + UserMessageId::new(), + [acp::ContentBlock::from("hello from the user")], + path_style, + cx, + ); + }); + }); + cx.run_until_parked(); + + // Reproduce the orphaned state from the bug: the sidebar metadata and + // serialized panel still reference the session, but the per-session + // async content save never landed, so the content row is absent. + let database = cx.update(|cx| ThreadsDatabase::connect(cx)).await.unwrap(); + database.delete_thread(session_id.clone()).await.unwrap(); + assert!( + database + .load_thread(session_id.clone()) + .await + .unwrap() + .is_none(), + "precondition: content row should be missing before the quit flush" + ); + + // Quitting must re-commit the content so the thread can be restored. + let flush = cx.update(|cx| agent.update(cx, |agent, cx| agent.flush_threads_on_quit(cx))); + flush.await; + + let restored = database + .load_thread(session_id.clone()) + .await + .unwrap() + .expect("thread content should be persisted to the database on quit"); + assert_eq!( + restored.messages.len(), + 1, + "the user message should survive the quit flush" + ); + } + #[test] fn test_ambiguous_mcp_prompt_names() { // Reserving the built-in `/compact` forces a same-named MCP prompt to be