fix(workspace): save app state before shutdown

app/src/integration_testing/persistence.rs

@@ -1,2 +1,19 @@
 #[cfg(feature = "local_fs")]
 pub use crate::persistence::database_file_path;
+
+use warpui::integration::TestStep;
+
+/// Replays the persistence portion of the app's `on_will_terminate` callback (see
+/// `app_callbacks` in `lib.rs`): enqueue a final session snapshot, then synchronously
+/// terminate the sqlite writer thread.
+///
+/// Integration tests cannot assert anything after the real termination callback runs, so
+/// this helper runs the same shutdown sequence mid-test, letting a later step verify that
+/// the snapshot reached the database.
+pub fn run_shutdown_persistence_hooks() -> TestStep {
+    TestStep::new("Run the shutdown persistence hooks").with_action(|app, _, _data| {
+        app.update(|ctx| {
+            crate::workspace::run_shutdown_persistence(ctx);
+        });
+    })
+}

app/src/lib.rs

@@ -2029,9 +2029,7 @@ pub(crate) fn app_callbacks(is_integration_test: bool) -> warpui::platform::AppC
                 manager.close_notebooks(ctx);
             });
 
-            PersistenceWriter::handle(ctx).update(ctx, |writer, _ctx| {
-                writer.terminate();
-            });
+            workspace::run_shutdown_persistence(ctx);
 
             // Shutdown all LSP servers gracefully before app termination
             lsp::LspManagerModel::handle(ctx).update(ctx, |manager, ctx| {

app/src/persistence/sqlite.rs

@@ -783,6 +783,36 @@ fn deduplicate_events(events: Vec<ModelEvent>) -> Vec<ModelEvent> {
     }
 }
 
+/// Deletes every persisted app-state row (windows, tabs, panes, panels, etc.).
+///
+/// Used by `save_app_state` before inserting a fresh snapshot, and by
+/// integration-test helpers that need to reset the persisted session.
+pub(super) fn delete_app_state(conn: &mut SqliteConnection) -> Result<(), Error> {
+    diesel::delete(schema::app::dsl::app).execute(conn)?;
+    diesel::delete(schema::terminal_panes::dsl::terminal_panes).execute(conn)?;
+    diesel::delete(schema::notebook_panes::dsl::notebook_panes).execute(conn)?;
+    diesel::delete(schema::code_panes::dsl::code_panes).execute(conn)?;
+    diesel::delete(schema::env_var_collection_panes::dsl::env_var_collection_panes)
+        .execute(conn)?;
+    diesel::delete(schema::workflow_panes::dsl::workflow_panes).execute(conn)?;
+    diesel::delete(schema::settings_panes::dsl::settings_panes).execute(conn)?;
+    diesel::delete(schema::ai_memory_panes::dsl::ai_memory_panes).execute(conn)?;
+    diesel::delete(schema::ai_document_panes::dsl::ai_document_panes).execute(conn)?;
+    diesel::delete(schema::mcp_server_panes::dsl::mcp_server_panes).execute(conn)?;
+    diesel::delete(schema::code_review_panes::dsl::code_review_panes).execute(conn)?;
+    diesel::delete(schema::ambient_agent_panes::dsl::ambient_agent_panes).execute(conn)?;
+    diesel::delete(schema::welcome_panes::dsl::welcome_panes).execute(conn)?;
+    diesel::delete(schema::browser_panes::dsl::browser_panes).execute(conn)?;
+    diesel::delete(schema::pane_leaves::dsl::pane_leaves).execute(conn)?;
+    diesel::delete(schema::pane_branches::dsl::pane_branches).execute(conn)?;
+    diesel::delete(schema::pane_nodes::dsl::pane_nodes).execute(conn)?;
+    diesel::delete(schema::tabs::dsl::tabs).execute(conn)?;
+    diesel::delete(schema::windows::dsl::windows).execute(conn)?;
+    diesel::delete(schema::active_mcp_servers::dsl::active_mcp_servers).execute(conn)?;
+    diesel::delete(schema::panels::dsl::panels).execute(conn)?;
+    Ok(())
+}
+
 // Used in the save_app_state function to help make the code more readable.
 struct SaveAppStateNodeTraversal<'a> {
     node: &'a PaneNodeSnapshot,
@@ -795,28 +825,7 @@ struct SaveAppStateNodeTraversal<'a> {
 fn save_app_state(conn: &mut SqliteConnection, app_state: &AppState) -> Result<()> {
     conn.transaction::<(), Error, _>(|conn| {
         // Remove old app state
-        diesel::delete(schema::app::dsl::app).execute(conn)?;
-        diesel::delete(schema::terminal_panes::dsl::terminal_panes).execute(conn)?;
-        diesel::delete(schema::notebook_panes::dsl::notebook_panes).execute(conn)?;
-        diesel::delete(schema::code_panes::dsl::code_panes).execute(conn)?;
-        diesel::delete(schema::env_var_collection_panes::dsl::env_var_collection_panes)
-            .execute(conn)?;
-        diesel::delete(schema::workflow_panes::dsl::workflow_panes).execute(conn)?;
-        diesel::delete(schema::settings_panes::dsl::settings_panes).execute(conn)?;
-        diesel::delete(schema::ai_memory_panes::dsl::ai_memory_panes).execute(conn)?;
-        diesel::delete(schema::ai_document_panes::dsl::ai_document_panes).execute(conn)?;
-        diesel::delete(schema::mcp_server_panes::dsl::mcp_server_panes).execute(conn)?;
-        diesel::delete(schema::code_review_panes::dsl::code_review_panes).execute(conn)?;
-        diesel::delete(schema::ambient_agent_panes::dsl::ambient_agent_panes).execute(conn)?;
-        diesel::delete(schema::welcome_panes::dsl::welcome_panes).execute(conn)?;
-        diesel::delete(schema::browser_panes::dsl::browser_panes).execute(conn)?;
-        diesel::delete(schema::pane_leaves::dsl::pane_leaves).execute(conn)?;
-        diesel::delete(schema::pane_branches::dsl::pane_branches).execute(conn)?;
-        diesel::delete(schema::pane_nodes::dsl::pane_nodes).execute(conn)?;
-        diesel::delete(schema::tabs::dsl::tabs).execute(conn)?;
-        diesel::delete(schema::windows::dsl::windows).execute(conn)?;
-        diesel::delete(schema::active_mcp_servers::dsl::active_mcp_servers).execute(conn)?;
-        diesel::delete(schema::panels::dsl::panels).execute(conn)?;
+        delete_app_state(conn)?;
 
         let mut active_window_id = None;
 

app/src/persistence/testing.rs

@@ -1,9 +1,11 @@
 //! Module with integration test-only util methods setting up sqlite.
 
-use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl};
-
-use super::{schema, sqlite::init_db};
+use diesel::{Connection, ExpressionMethods, QueryDsl, RunQueryDsl};
 
+use super::{
+    schema,
+    sqlite::{database_file_path, establish_ro_connection, init_db},
+};
 /// Updates the 'user' and 'host' columns for stored blocks to the given values.
 ///
 /// This is used at runtime to update the user and host values to real values based on the running
@@ -44,3 +46,33 @@ pub fn set_user_and_hostname_for_commands(user: String, hostname: String) {
     .execute(&mut conn)
     .expect("Failed to update user and hostname for persisted commands.");
 }
+
+/// Returns the number of tabs stored in the persisted app-state snapshot.
+///
+/// This is used by integration tests to verify that a session snapshot actually reached the
+/// sqlite database (e.g. via the shutdown save hook).
+pub fn count_persisted_tabs() -> i64 {
+    let database_path = database_file_path();
+    let database_url = database_path
+        .to_str()
+        .expect("SQLite database path should be valid UTF-8.");
+    let mut conn = establish_ro_connection(database_url)
+        .expect("Should be able to establish read-only sqlite connection.");
+
+    schema::tabs::dsl::tabs
+        .count()
+        .get_result(&mut conn)
+        .expect("Failed to count persisted tabs.")
+}
+
+/// Deletes the persisted app-state snapshot (windows, tabs, panes, etc.).
+///
+/// This lets integration tests wipe state written by ambient saves (window
+/// events, tab actions) so they can verify that a later save — e.g. the
+/// shutdown hook — persists a fresh snapshot on its own.
+pub fn clear_persisted_app_state() {
+    let mut conn = init_db().expect("Should be able to establish sqlite connection.");
+
+    conn.transaction(super::sqlite::delete_app_state)
+        .expect("Failed to clear persisted app state.");
+}

app/src/workspace/global_actions.rs

@@ -11,6 +11,7 @@ use warp_core::execution_mode::AppExecutionMode;
 
 use crate::ai::agent::conversation::AIConversationId;
 use crate::ai::agent::AIAgentExchangeId;
+use crate::persistence::PersistenceWriter;
 use crate::root_view::OpenPath;
 use crate::undo_close::UndoCloseStack;
 use crate::workspace::{Workspace, WorkspaceAction};
@@ -157,6 +158,19 @@ fn save_app(_: &(), ctx: &mut AppContext) {
     }
 }
 
+/// Persists a final session snapshot and then shuts down the sqlite writer.
+///
+/// Called from the app's `on_will_terminate` callback (see `app_callbacks` in `lib.rs`).
+/// The snapshot must be enqueued before `PersistenceWriter::terminate`, which synchronously
+/// drains the writer queue, so the session state at quit reaches the database.
+pub(crate) fn run_shutdown_persistence(ctx: &mut AppContext) {
+    save_app(&(), ctx);
+
+    PersistenceWriter::handle(ctx).update(ctx, |writer, _ctx| {
+        writer.terminate();
+    });
+}
+
 fn toggle_debug_network_status(_: &(), ctx: &mut AppContext) {
     NetworkStatus::handle(ctx).update(ctx, move |me, ctx| {
         let is_reachable = me.is_online();
@@ -249,3 +263,68 @@ fn summarize_ai_conversation(prompt: &Option<String>, ctx: &mut AppContext) {
 fn trigger_log_out(_: &(), ctx: &mut AppContext) {
     auth::log_out(ctx)
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::mpsc::sync_channel;
+
+    use warpui::App;
+
+    use crate::{
+        persistence::{ModelEvent, PersistenceWriter, WriterHandles},
+        test_util::settings::initialize_settings_for_tests,
+        workspace::{cross_window_tab_drag::CrossWindowTabDrag, WorkspaceRegistry},
+        GlobalResourceHandles, GlobalResourceHandlesProvider,
+    };
+
+    use super::*;
+
+    #[test]
+    fn shutdown_save_sends_snapshot_before_writer_termination() {
+        App::test((), |mut app| async move {
+            initialize_settings_for_tests(&mut app);
+            app.add_singleton_model(|_| WorkspaceRegistry::new());
+            app.add_singleton_model(|_| CrossWindowTabDrag::new());
+
+            // Mirror the production wiring in `sqlite::start_writer`: `save_app`
+            // and `PersistenceWriter::terminate` share one channel into the
+            // writer thread. Capacity 2 fits the Snapshot + Terminate events;
+            // the dummy thread stands in for the writer so `terminate` has a
+            // real handle to join.
+            let (tx, rx) = sync_channel(2);
+            let writer_sender = tx.clone();
+            let thread_handle = std::thread::spawn(|| {});
+            app.add_singleton_model(move |_| {
+                PersistenceWriter::new(Some(WriterHandles {
+                    handle: thread_handle,
+                    sender: writer_sender,
+                }))
+            });
+
+            let mut global_resource_handles = GlobalResourceHandles::mock(&mut app);
+            global_resource_handles.model_event_sender = Some(tx);
+            app.add_singleton_model(|_| {
+                GlobalResourceHandlesProvider::new(global_resource_handles)
+            });
+
+            app.update(|ctx| {
+                run_shutdown_persistence(ctx);
+            });
+
+            let first = rx
+                .try_recv()
+                .expect("shutdown save should enqueue a snapshot");
+            assert!(
+                matches!(first, ModelEvent::Snapshot(_)),
+                "the session snapshot should be enqueued first, got {first:?}"
+            );
+            let second = rx
+                .try_recv()
+                .expect("shutdown save should terminate the writer");
+            assert!(
+                matches!(second, ModelEvent::Terminate),
+                "writer termination should follow the snapshot, got {second:?}"
+            );
+        });
+    }
+}

app/src/workspace/mod.rs

@@ -53,6 +53,7 @@ pub use action::{
     VerticalTabsPaneContextMenuTarget, WorkspaceAction,
 };
 pub use active_session::ActiveSession;
+pub(crate) use global_actions::run_shutdown_persistence;
 pub use global_actions::{
     ForkAIConversationParams, ForkFromExchange, ForkedConversationDestination,
 };

crates/integration/src/bin/integration.rs

@@ -187,6 +187,7 @@ fn register_tests() -> HashMap<&'static str, BoxedBuilderFn> {
     register_test!(test_disabling_action_dispatching);
     register_test!(test_session_restoration);
     register_test!(test_restored_blocks_on_different_hosts);
+    register_test!(test_shutdown_save_persists_session_snapshot);
     register_test!(test_restore_snapshot_with_deleted_cwd);
     register_test!(test_session_restoration_with_multiple_shells);
     register_test!(test_restore_snapshot_with_background_output);

crates/integration/src/test/session_restoration.rs

@@ -12,16 +12,19 @@ use warp::{
         terminal::wait_until_bootstrapped_single_pane_for_tab,
         view_getters::single_terminal_view_for_tab,
         workflow::assert_workflow_metadata_revision,
+        workspace::assert_tab_count,
     },
     settings::Preference,
     settings_view::{SettingsSection, SettingsView},
-    sqlite_testing::set_user_and_hostname_for_blocks,
+    sqlite_testing::{
+        clear_persisted_app_state, count_persisted_tabs, set_user_and_hostname_for_blocks,
+    },
     terminal::{
         model::{session::get_local_hostname, terminal_model::BlockIndex},
         shell::ShellType,
         History, ShellHost, TerminalView,
     },
-    workspace::Workspace,
+    workspace::{Workspace, NEW_TAB_BUTTON_POSITION_ID},
 };
 use warpui::{
     async_assert_eq,
@@ -546,3 +549,59 @@ pub fn test_restore_snapshot_with_settings_page() -> Builder {
                 }),
         )
 }
+
+/// Tests that the shutdown save hook persists the live session before the sqlite
+/// writer terminates.
+///
+/// The app's `on_will_terminate` callback must enqueue a final session snapshot
+/// before terminating the persistence writer, so that session restoration reflects
+/// the state at quit rather than the last ambient save. This test opens a second
+/// tab, waits for the ambient save triggered by that action to flush, wipes the
+/// persisted snapshot, and then replays the shutdown persistence hooks — so the
+/// snapshot found afterwards can only have been written by the shutdown save.
+pub fn test_shutdown_save_persists_session_snapshot() -> Builder {
+    new_builder()
+        .with_step(wait_until_bootstrapped_single_pane_for_tab(0))
+        .with_step(
+            new_step_with_default_assertions("Add a second tab with the new tab button")
+                .with_click_on_saved_position(NEW_TAB_BUTTON_POSITION_ID)
+                .add_assertion(assert_tab_count(2)),
+        )
+        .with_step(wait_until_bootstrapped_single_pane_for_tab(1))
+        .with_step(
+            TestStep::new("Wait for the ambient save to flush both tabs").add_named_assertion(
+                "Ambient save persisted both tabs",
+                |_app, _window_id| {
+                    async_assert_eq!(
+                        count_persisted_tabs(),
+                        2,
+                        "The tab action's ambient save should persist both tabs"
+                    )
+                },
+            ),
+        )
+        .with_step(
+            TestStep::new("Clear the persisted snapshot to isolate the shutdown save")
+                .with_action(|_app, _window_id, _data| clear_persisted_app_state())
+                .add_named_assertion("Persisted snapshot is empty", |_app, _window_id| {
+                    async_assert_eq!(
+                        count_persisted_tabs(),
+                        0,
+                        "Clearing the persisted app state should remove all tabs"
+                    )
+                }),
+        )
+        .with_step(integration_testing::persistence::run_shutdown_persistence_hooks())
+        .with_step(
+            TestStep::new("Assert the persisted snapshot contains both tabs").add_named_assertion(
+                "Persisted snapshot contains both tabs",
+                |_app, _window_id| {
+                    async_assert_eq!(
+                        count_persisted_tabs(),
+                        2,
+                        "The shutdown save should persist a snapshot with both tabs"
+                    )
+                },
+            ),
+        )
+}

crates/integration/tests/integration/ui_tests.rs

@@ -58,6 +58,7 @@ integration_tests! {
     test_disabling_action_dispatching,
     test_session_restoration,
     test_restored_blocks_on_different_hosts,
+    test_shutdown_save_persists_session_snapshot,
     test_restore_snapshot_with_deleted_cwd,
     test_session_restoration_with_multiple_shells,
     test_restore_snapshot_with_background_output,