Fix orphan cleanup to properly handle workspaces containing multiple worktrees (Vibe Kanban) (#1545)

LSRCT · web-flow · commit a8cc7746cd60 · 2025-12-15T19:11:12.000Z
* All done. Here's a summary of the changes made:

Fixed the orphan cleanup bug by moving the logic from `LocalContainerService` to `WorkspaceManager` and properly handling workspaces (which contain multiple worktrees).

1. **`crates/services/src/services/workspace_manager.rs`**:
   - Added `cleanup_orphan_workspaces()` - public method that scans the workspace base directory for orphaned workspaces and cleans them up
   - Added `cleanup_workspace_without_repos()` - helper method that cleans up a workspace by discovering worktrees (subdirectories with `.git` files) and cleaning each one properly using `WorktreeManager::cleanup_worktree()`

2. **`crates/local-deployment/src/container.rs`**:
   - Removed the broken `cleanup_orphaned_worktrees()` method (was incorrectly treating workspaces as worktrees)
   - Updated `spawn_workspace_cleanup()` to call `WorkspaceManager::cleanup_orphan_workspaces()` instead
   - Removed unused imports (`WorktreeCleanup`, `WorktreeManager`)

The old code was passing workspace paths (e.g., `/worktrees/abc-task-title/`) directly to `WorktreeManager::cleanup_worktree()`, which expected individual worktree paths (e.g., `/worktrees/abc-task-title/repo-name/`). The new code correctly:
1. Identifies orphaned **workspaces** by checking if the workspace path exists in any task attempt's `container_ref`
2. For each orphaned workspace, scans its subdirectories for worktrees (directories with `.git` files)
3. Cleans up each worktree properly using `WorktreeManager::cleanup_worktree()`
4. Finally removes the workspace directory itself

* Done. I've added `cleanup_suspected_worktree()` to `WorktreeManager` and updated `cleanup_workspace_without_repos()` in `WorkspaceManager` to use it. This keeps the worktree detection and cleanup logic in the appropriate manager.

* Comments

* lint
diff --git a/crates/local-deployment/src/container.rs b/crates/local-deployment/src/container.rs
@@ -52,7 +52,6 @@ use services::services::{
     queued_message::QueuedMessageService,
     share::SharePublisher,
     workspace_manager::{RepoWorkspaceInput, WorkspaceManager},
-    worktree_manager::{WorktreeCleanup, WorktreeManager},
 };
 use tokio::{sync::RwLock, task::JoinHandle};
 use tokio_util::io::ReaderStream;
@@ -143,73 +142,7 @@ impl LocalContainerService {
         map.remove(id)
     }
 
-    /// Find and delete orphaned worktrees that don't correspond to any task attempts
-    async fn cleanup_orphaned_worktrees(&self) {
-        // Check if orphan cleanup is disabled via environment variable
-        if std::env::var("DISABLE_WORKTREE_ORPHAN_CLEANUP").is_ok() {
-            tracing::debug!(
-                "Orphan worktree cleanup is disabled via DISABLE_WORKTREE_ORPHAN_CLEANUP environment variable"
-            );
-            return;
-        }
-        let worktree_base_dir = WorktreeManager::get_worktree_base_dir();
-        if !worktree_base_dir.exists() {
-            tracing::debug!(
-                "Worktree base directory {} does not exist, skipping orphan cleanup",
-                worktree_base_dir.display()
-            );
-            return;
-        }
-        let entries = match std::fs::read_dir(&worktree_base_dir) {
-            Ok(entries) => entries,
-            Err(e) => {
-                tracing::error!(
-                    "Failed to read worktree base directory {}: {}",
-                    worktree_base_dir.display(),
-                    e
-                );
-                return;
-            }
-        };
-        for entry in entries {
-            let entry = match entry {
-                Ok(entry) => entry,
-                Err(e) => {
-                    tracing::warn!("Failed to read directory entry: {}", e);
-                    continue;
-                }
-            };
-            let path = entry.path();
-            // Only process directories
-            if !path.is_dir() {
-                continue;
-            }
-
-            let worktree_path_str = path.to_string_lossy().to_string();
-            if let Ok(false) =
-                TaskAttempt::container_ref_exists(&self.db().pool, &worktree_path_str).await
-            {
-                // This is an orphaned worktree - delete it
-                tracing::info!("Found orphaned worktree: {}", worktree_path_str);
-                if let Err(e) =
-                    WorktreeManager::cleanup_worktree(&WorktreeCleanup::new(path, None)).await
-                {
-                    tracing::error!(
-                        "Failed to remove orphaned worktree {}: {}",
-                        worktree_path_str,
-                        e
-                    );
-                } else {
-                    tracing::info!(
-                        "Successfully removed orphaned worktree: {}",
-                        worktree_path_str
-                    );
-                }
-            }
-        }
-    }
-
-    async fn cleanup_attempt_workspace(db: &DBService, attempt: &TaskAttempt) {
+    pub async fn cleanup_attempt_workspace(db: &DBService, attempt: &TaskAttempt) {
         let Some(container_ref) = &attempt.container_ref else {
             return;
         };
@@ -224,10 +157,10 @@ impl LocalContainerService {
                 "No repositories found for attempt {}, cleaning up workspace directory only",
                 attempt.id
             );
-            if workspace_dir.exists() {
-                if let Err(e) = tokio::fs::remove_dir_all(&workspace_dir).await {
-                    tracing::warn!("Failed to remove workspace directory: {}", e);
-                }
+            if workspace_dir.exists()
+                && let Err(e) = tokio::fs::remove_dir_all(&workspace_dir).await
+            {
+                tracing::warn!("Failed to remove workspace directory: {}", e);
             }
         } else {
             WorkspaceManager::cleanup_workspace(&workspace_dir, &repositories)
@@ -264,7 +197,7 @@ impl LocalContainerService {
     pub async fn spawn_workspace_cleanup(&self) {
         let db = self.db.clone();
         let mut cleanup_interval = tokio::time::interval(tokio::time::Duration::from_secs(1800)); // 30 minutes
-        self.cleanup_orphaned_worktrees().await;
+        WorkspaceManager::cleanup_orphan_workspaces(&self.db.pool).await;
         tokio::spawn(async move {
             loop {
                 cleanup_interval.tick().await;
@@ -778,19 +711,19 @@ impl LocalContainerService {
         let repos = AttemptRepo::find_repos_with_copy_files(&self.db.pool, task_attempt.id).await?;
 
         for repo in &repos {
-            if let Some(copy_files) = &repo.copy_files {
-                if !copy_files.trim().is_empty() {
-                    let worktree_path = workspace_dir.join(&repo.name);
-                    self.copy_project_files(&repo.path, &worktree_path, copy_files)
-                        .await
-                        .unwrap_or_else(|e| {
-                            tracing::warn!(
-                                "Failed to copy project files for repo '{}': {}",
-                                repo.name,
-                                e
-                            );
-                        });
-                }
+            if let Some(copy_files) = &repo.copy_files
+                && !copy_files.trim().is_empty()
+            {
+                let worktree_path = workspace_dir.join(&repo.name);
+                self.copy_project_files(&repo.path, &worktree_path, copy_files)
+                    .await
+                    .unwrap_or_else(|e| {
+                        tracing::warn!(
+                            "Failed to copy project files for repo '{}': {}",
+                            repo.name,
+                            e
+                        );
+                    });
             }
         }
 
diff --git a/crates/local-deployment/src/copy.rs b/crates/local-deployment/src/copy.rs
@@ -101,12 +101,11 @@ fn copy_single_file(
     let target_file = target_root.join(relative_path);
 
     // Skip if target exists with same size
-    if let Ok(target_meta) = fs::metadata(&target_file) {
-        if let Ok(source_meta) = fs::metadata(source_file) {
-            if target_meta.len() == source_meta.len() {
-                return Ok(false);
-            }
-        }
+    if let Ok(target_meta) = fs::metadata(&target_file)
+        && let Ok(source_meta) = fs::metadata(source_file)
+        && target_meta.len() == source_meta.len()
+    {
+        return Ok(false);
     }
 
     if let Some(parent) = target_file.parent()
diff --git a/crates/services/src/services/workspace_manager.rs b/crates/services/src/services/workspace_manager.rs
@@ -1,8 +1,9 @@
 use std::path::{Path, PathBuf};
 
-use db::models::repo::Repo;
+use db::models::{repo::Repo, task_attempt::TaskAttempt};
+use sqlx::{Pool, Sqlite};
 use thiserror::Error;
-use tracing::{debug, error, info};
+use tracing::{debug, error, info, warn};
 use uuid::Uuid;
 
 use super::worktree_manager::{WorktreeCleanup, WorktreeError, WorktreeManager};
@@ -282,4 +283,107 @@ impl WorkspaceManager {
             }
         }
     }
+
+    pub async fn cleanup_orphan_workspaces(db: &Pool<Sqlite>) {
+        if std::env::var("DISABLE_WORKTREE_ORPHAN_CLEANUP").is_ok() {
+            debug!(
+                "Orphan workspace cleanup is disabled via DISABLE_WORKTREE_ORPHAN_CLEANUP environment variable"
+            );
+            return;
+        }
+
+        let workspace_base_dir = Self::get_workspace_base_dir();
+        if !workspace_base_dir.exists() {
+            debug!(
+                "Workspace base directory {} does not exist, skipping orphan cleanup",
+                workspace_base_dir.display()
+            );
+            return;
+        }
+
+        let entries = match std::fs::read_dir(&workspace_base_dir) {
+            Ok(entries) => entries,
+            Err(e) => {
+                error!(
+                    "Failed to read workspace base directory {}: {}",
+                    workspace_base_dir.display(),
+                    e
+                );
+                return;
+            }
+        };
+
+        for entry in entries {
+            let entry = match entry {
+                Ok(entry) => entry,
+                Err(e) => {
+                    warn!("Failed to read directory entry: {}", e);
+                    continue;
+                }
+            };
+
+            let path = entry.path();
+            if !path.is_dir() {
+                continue;
+            }
+
+            let workspace_path_str = path.to_string_lossy().to_string();
+            if let Ok(false) = TaskAttempt::container_ref_exists(db, &workspace_path_str).await {
+                info!("Found orphaned workspace: {}", workspace_path_str);
+                if let Err(e) = Self::cleanup_workspace_without_repos(&path).await {
+                    error!(
+                        "Failed to remove orphaned workspace {}: {}",
+                        workspace_path_str, e
+                    );
+                } else {
+                    info!(
+                        "Successfully removed orphaned workspace: {}",
+                        workspace_path_str
+                    );
+                }
+            }
+        }
+    }
+
+    async fn cleanup_workspace_without_repos(workspace_dir: &Path) -> Result<(), WorkspaceError> {
+        info!(
+            "Cleaning up orphaned workspace at {}",
+            workspace_dir.display()
+        );
+
+        let entries = match std::fs::read_dir(workspace_dir) {
+            Ok(entries) => entries,
+            Err(e) => {
+                debug!(
+                    "Cannot read workspace directory {}, attempting direct removal: {}",
+                    workspace_dir.display(),
+                    e
+                );
+                return tokio::fs::remove_dir_all(workspace_dir)
+                    .await
+                    .map_err(WorkspaceError::Io);
+            }
+        };
+
+        for entry in entries.filter_map(|e| e.ok()) {
+            let path = entry.path();
+            if path.is_dir()
+                && let Err(e) = WorktreeManager::cleanup_suspected_worktree(&path).await
+            {
+                warn!("Failed to cleanup suspected worktree: {}", e);
+            }
+        }
+
+        if workspace_dir.exists()
+            && let Err(e) = tokio::fs::remove_dir_all(workspace_dir).await
+        {
+            debug!(
+                "Could not remove workspace directory {}: {}",
+                workspace_dir.display(),
+                e
+            );
+        }
+
+        Ok(())
+    }
 }
diff --git a/crates/services/src/services/worktree_manager.rs b/crates/services/src/services/worktree_manager.rs
@@ -518,4 +518,16 @@ impl WorktreeManager {
     pub fn get_worktree_base_dir() -> std::path::PathBuf {
         utils::path::get_vibe_kanban_temp_dir().join("worktrees")
     }
+
+    pub async fn cleanup_suspected_worktree(path: &Path) -> Result<bool, WorktreeError> {
+        let git_marker = path.join(".git");
+        if !git_marker.exists() || !git_marker.is_file() {
+            return Ok(false);
+        }
+
+        debug!("Cleaning up suspected worktree at {}", path.display());
+        let cleanup = WorktreeCleanup::new(path.to_path_buf(), None);
+        Self::cleanup_worktree(&cleanup).await?;
+        Ok(true)
+    }
 }
