From bd2aafd7b842cfefbf8f71929e21295c4fe61c70 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EA=B9=80=EC=9A=A9=EC=84=B1?= <66245186+kys0213@users.noreply.github.com> Date: Mon, 30 Mar 2026 16:39:11 +0900 Subject: [PATCH] fix(daemon): ensure rollback persists items to DB for worktree continuity Items collected from DataSource live in-memory only and may not exist in the DB when rollback_running_to_pending() is called during graceful shutdown. The UPDATE call would silently fail with ItemNotFound, losing the previous_worktree_path on restart. Now the rollback ensures the item row exists (via insert) before updating its worktree state. Closes #679 Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/belt-daemon/src/daemon.rs | 117 ++++++++++++++++++++++++++++--- 1 file changed, 108 insertions(+), 9 deletions(-) diff --git a/crates/belt-daemon/src/daemon.rs b/crates/belt-daemon/src/daemon.rs index 35ad830..1f3e412 100644 --- a/crates/belt-daemon/src/daemon.rs +++ b/crates/belt-daemon/src/daemon.rs @@ -2087,19 +2087,30 @@ impl Daemon { } // Persist worktree state to DB so it survives restart. - if let Some(db) = &self.db - && let Err(e) = db.update_item_worktree_state( + // Items collected from DataSource live in-memory only and may + // not yet exist in the DB. Ensure the row exists before updating. + if let Some(db) = &self.db { + if db.get_item(&item.work_id).is_err() + && let Err(e) = db.insert_item(item) + { + tracing::warn!( + work_id = %item.work_id, + error = %e, + "failed to insert item into DB during rollback" + ); + } + if let Err(e) = db.update_item_worktree_state( &item.work_id, QueuePhase::Pending, item.worktree_preserved, item.previous_worktree_path.as_deref(), - ) - { - tracing::warn!( - work_id = %item.work_id, - error = %e, - "failed to persist worktree state to DB during rollback" - ); + ) { + tracing::warn!( + work_id = %item.work_id, + error = %e, + "failed to persist worktree state to DB during rollback" + ); + } } self.tracker.release(&ws_name); @@ -4692,6 +4703,94 @@ sources: } } + // --------------------------------------------------------------- + // Worktree continuity across graceful shutdown (issue #679) + // --------------------------------------------------------------- + + #[test] + fn rollback_inserts_item_into_db_when_not_present() { + // Items collected from DataSource live in-memory only. + // rollback_running_to_pending must insert the item into the DB + // before updating its worktree state so the path survives restart. + let tmp = TempDir::new().unwrap(); + let source = MockDataSource::new("github"); + let daemon = setup_daemon(&tmp, source, vec![]); + + let db = Database::open_in_memory().unwrap(); + let mut daemon = daemon.with_db(db); + + // Create the worktree directory so rollback registers it. + let ws_path = daemon.worktree_mgr.path("test-ws"); + std::fs::create_dir_all(&ws_path).unwrap(); + + let mut item = test_item("github:org/repo#99", "analyze"); + item.phase = QueuePhase::Running; + daemon.push_item(item); + + // Item is NOT in the DB yet (only in-memory queue). + assert!( + daemon + .db + .as_ref() + .unwrap() + .get_item("github:org/repo#99:analyze") + .is_err() + ); + + daemon.rollback_running_to_pending(); + + // After rollback, item should exist in DB with previous_worktree_path. + let db_item = daemon + .db + .as_ref() + .unwrap() + .get_item("github:org/repo#99:analyze") + .expect("item should be inserted into DB during rollback"); + assert_eq!(db_item.phase, QueuePhase::Pending); + assert!(db_item.worktree_preserved); + assert_eq!( + db_item.previous_worktree_path.as_deref(), + Some(ws_path.to_str().unwrap()) + ); + } + + #[test] + fn rollback_updates_existing_db_item_worktree_state() { + // When the item already exists in the DB (e.g., from a cron job), + // rollback should update its worktree state without inserting a + // duplicate. + let tmp = TempDir::new().unwrap(); + let source = MockDataSource::new("github"); + let daemon = setup_daemon(&tmp, source, vec![]); + + let db = Database::open_in_memory().unwrap(); + // Pre-insert the item into DB. + let pre_item = test_item("github:org/repo#100", "analyze"); + db.insert_item(&pre_item).unwrap(); + let mut daemon = daemon.with_db(db); + + // Create the worktree directory so rollback registers it. + let ws_path = daemon.worktree_mgr.path("test-ws"); + std::fs::create_dir_all(&ws_path).unwrap(); + + let mut item = test_item("github:org/repo#100", "analyze"); + item.phase = QueuePhase::Running; + daemon.push_item(item); + + daemon.rollback_running_to_pending(); + + // DB should have the updated worktree state. + let db_item = daemon + .db + .as_ref() + .unwrap() + .get_item("github:org/repo#100:analyze") + .unwrap(); + assert_eq!(db_item.phase, QueuePhase::Pending); + assert!(db_item.worktree_preserved); + assert!(db_item.previous_worktree_path.is_some()); + } + // --------------------------------------------------------------- // report_dir initialization tests // ---------------------------------------------------------------